def testScanAndDumpPopulatesMemoryRegions(self): client_mock = action_mocks.MultiGetFileClientMock( memory_actions.YaraProcessScan, memory_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles) procs = [p for p in self.procs if p.pid in [108]] with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid))): session_id = flow_test_lib.TestFlowHelper( memory.YaraProcessScan.__name__, client_mock, yara_signature=_TEST_YARA_SIGNATURE, client_id=self.client_id, token=self.token, include_errors_in_results="ALL_ERRORS", include_misses_in_results=True, dump_process_on_match=True) results = flow_test_lib.GetFlowResults(self.client_id, session_id) dumps = [ r for r in results if isinstance(r, rdf_memory.YaraProcessDumpResponse) ] self.assertLen(dumps, 1) self.assertLen(dumps[0].dumped_processes, 1) self.assertLen(dumps[0].dumped_processes[0].memory_regions, 2) regions = dumps[0].dumped_processes[0].memory_regions self.assertEqual(regions[0].start, 0) self.assertEqual(regions[0].size, 100) self.assertEqual(regions[0].dumped_size, 100) self.assertEqual(regions[0].is_executable, True) self.assertEqual(regions[0].is_writable, True) self.assertIsNotNone(regions[0].file) self.assertEqual(regions[1].start, 1000) self.assertEqual(regions[1].size, 104) self.assertEqual(regions[1].dumped_size, 104) self.assertEqual(regions[1].is_executable, False) self.assertEqual(regions[1].is_writable, False) self.assertIsNotNone(regions[1].file)
def testScanAndDump(self): client_mock = action_mocks.MultiGetFileClientMock( memory_actions.YaraProcessScan, memory_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles) procs = [p for p in self.procs if p.pid in [102, 103]] with mock.patch.object(file_store.EXTERNAL_FILE_STORE, "AddFiles") as efs: with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid))): session_id = flow_test_lib.TestFlowHelper( memory.YaraProcessScan.__name__, client_mock, yara_signature=test_yara_signature, client_id=self.client_id, token=self.token, include_errors_in_results=True, include_misses_in_results=True, dump_process_on_match=True) # Process dumps are not pushed to external file stores. self.assertEqual(efs.call_count, 0) results = flow_test_lib.GetFlowResults(self.client_id.Basename(), session_id) # 1. Scan result match. # 2. Scan result miss. # 3. ProcDump response. # 4. Stat entry for the dumped file. self.assertLen(results, 4) self.assertIsInstance(results[0], rdf_memory.YaraProcessScanMatch) self.assertIsInstance(results[1], rdf_memory.YaraProcessScanMiss) self.assertIsInstance(results[2], rdf_memory.YaraProcessDumpResponse) self.assertIsInstance(results[3], rdf_client_fs.StatEntry) self.assertLen(results[2].dumped_processes, 1) self.assertEqual(results[0].process.pid, results[2].dumped_processes[0].process.pid) self.assertIn(str(results[2].dumped_processes[0].process.pid), results[3].pathspec.path)
def setUp(self): # pylint: disable=invalid-name """Set up test method.""" super(ApiRegressionTest, self).setUp() if not self.__class__.api_method: raise ValueError("%s.api_method has to be set." % self.__class__.__name__) if not self.__class__.handler: raise ValueError("%s.handler has to be set." % self.__class__.__name__) self.checks = [] p = psutil.Process(os.getpid()) self.syscalls_stubber = utils.MultiStubber( (socket, "gethostname", lambda: "test.host"), (os, "getpid", lambda: 42), (psutil, "Process", lambda _=None: p)) self.syscalls_stubber.Start() self.token.username = "******" webauth.WEBAUTH_MANAGER.SetUserName(self.token.username) # Force creation of new APIAuthorizationManager. api_auth_manager.APIACLInit.InitApiAuthManager() self.config_overrider = test_lib.ConfigOverrider({ # For regression tests we want to use a fixed version number instead of # current one. Otherwise regression data would have to be re-generated # each time GRR version is increased. "Source.version_major": 1, "Source.version_minor": 2, "Source.version_revision": 3, "Source.version_release": 4, "Source.version_string": "1.2.3.4", "Source.version_numeric": 1234, }) self.config_overrider.Start()
def testScanAndDumpPrioritizesRegionsWithMatch(self): client_mock = action_mocks.MultiGetFileClientMock( memory_actions.YaraProcessScan, memory_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles) procs = [p for p in self.procs if p.pid in [109]] with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid, tmp_dir=self._tmp_dir))): session_id = flow_test_lib.TestFlowHelper( memory.YaraProcessScan.__name__, client_mock, yara_signature=_TEST_YARA_SIGNATURE, client_id=self.client_id, creator=self.test_username, include_errors_in_results="ALL_ERRORS", include_misses_in_results=True, dump_process_on_match=True, process_dump_size_limit=100 + 104) # size of first and third region. results = flow_test_lib.GetFlowResults(self.client_id, session_id) dumps = [ r for r in results if isinstance(r, rdf_memory.YaraProcessDumpResponse) ] self.assertLen(dumps, 1) self.assertLen(dumps[0].dumped_processes, 1) self.assertLen(dumps[0].dumped_processes[0].memory_regions, 2) regions = dumps[0].dumped_processes[0].memory_regions # Dump should skip the second region, because the first and third fill the # size limit. self.assertEqual(regions[0].start, 0) self.assertEqual(regions[0].dumped_size, 100) self.assertIsNotNone(regions[0].file) self.assertEqual(regions[1].start, 101) self.assertEqual(regions[1].dumped_size, 104) self.assertIsNotNone(regions[1].file)
def ValidateConfig(self, config_file=None): """Iterate over all the sections in the config file and validate them.""" logging.debug("Processing %s", config_file) if isinstance(config_file, config_lib.GrrConfigManager): conf_obj = config_file else: conf_obj = config.CONFIG.MakeNewConfig() conf_obj.Initialize(filename=config_file, reset=True) with utils.MultiStubber((config, "CONFIG", conf_obj), (config_lib, "_CONFIG", conf_obj)): all_sections = conf_obj.GetSections() # TODO: Validate causes imports, which then register metrics, # which fails because StatsCollector has been initialized. with mock.patch.multiple(metrics, _metadata=[], _finalized=False): errors = conf_obj.Validate(sections=all_sections) return errors
def testCPULimit(self): received_messages = [] class MockWorker(object): def Heartbeat(self): pass def SendClientAlert(self, msg): received_messages.append(msg) class FakeProcess(object): times = [(1, 0), (2, 0), (3, 0), (10000, 0), (10001, 0)] def __init__(self, unused_pid=None): self.pcputimes = collections.namedtuple( "pcputimes", ["user", "system"]) def cpu_times(self): # pylint: disable=g-bad-name return self.pcputimes(*self.times.pop(0)) results = [] def MockSendReply(unused_self, reply=None, **kwargs): results.append(reply or rdf_client.LogMessage(**kwargs)) message = rdf_flows.GrrMessage(name="ProgressAction", cpu_limit=3600) action_cls = actions.ActionPlugin.classes[message.name] with utils.MultiStubber((psutil, "Process", FakeProcess), (action_cls, "SendReply", MockSendReply)): action_cls._authentication_required = False action = action_cls(grr_worker=MockWorker()) action.Execute(message) self.assertIn("Action exceeded cpu limit.", results[0].error_message) self.assertIn("CPUExceededError", results[0].error_message) self.assertLen(received_messages, 1) self.assertEqual(received_messages[0], "Cpu limit exceeded.")
def _RunYaraProcessScan(self, procs, action_mock=None, ignore_grr_process=False, include_errors_in_results="NO_ERRORS", include_misses_in_results=False, max_results_per_process=0, **kw): if action_mock is None: client_mock = action_mocks.ActionMock( memory_actions.YaraProcessScan) else: client_mock = action_mock with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid, tmp_dir=self._tmp_dir))): session_id = flow_test_lib.TestFlowHelper( memory.YaraProcessScan.__name__, client_mock, yara_signature=_TEST_YARA_SIGNATURE, client_id=self.client_id, ignore_grr_process=ignore_grr_process, include_errors_in_results=include_errors_in_results, include_misses_in_results=include_misses_in_results, max_results_per_process=max_results_per_process, creator=self.test_username, **kw) res = flow_test_lib.GetFlowResults(self.client_id, session_id) matches = [ r for r in res if isinstance(r, rdf_memory.YaraProcessScanMatch) ] errors = [ r for r in res if isinstance(r, rdf_memory.ProcessMemoryError) ] misses = [ r for r in res if isinstance(r, rdf_memory.YaraProcessScanMiss) ] return matches, errors, misses
def _RunProcessDump(self, pids=None, size_limit=None, chunk_size=None): procs = self.procs with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid, tmp_dir=self._tmp_dir))): client_mock = action_mocks.MultiGetFileClientMock( memory_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles) session_id = flow_test_lib.TestFlowHelper( memory.DumpProcessMemory.__name__, client_mock, pids=pids or [105], size_limit=size_limit, chunk_size=chunk_size, client_id=self.client_id, ignore_grr_process=True, creator=self.test_username) return flow_test_lib.GetFlowResults(self.client_id, session_id)
def _RunYaraProcessScan(self, procs, ignore_grr_process=False, include_errors_in_results=False, include_misses_in_results=False, max_results_per_process=0, **kw): client_mock = action_mocks.ActionMock(yara_actions.YaraProcessScan) with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid))): session_id = flow_test_lib.TestFlowHelper( yara_flows.YaraProcessScan.__name__, client_mock, yara_signature=test_yara_signature, client_id=self.client_id, ignore_grr_process=ignore_grr_process, include_errors_in_results=include_errors_in_results, include_misses_in_results=include_misses_in_results, max_results_per_process=max_results_per_process, token=self.token, **kw) flow_obj = aff4.FACTORY.Open(session_id) results = flow_obj.TypedResultCollection() matches = [ x[1].payload for x in results.ScanByType(rdf_yara.YaraProcessScanMatch.__name__) ] errors = [ x[1].payload for x in results.ScanByType(rdf_yara.YaraProcessError.__name__) ] misses = [ x[1].payload for x in results.ScanByType(rdf_yara.YaraProcessScanMiss.__name__) ] return (matches, errors, misses)
def _RunProcessDump(self, pids=None, size_limit=None, chunk_size=None): procs = self.procs with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid))): client_mock = action_mocks.MultiGetFileClientMock( yara_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles) session_id = flow_test_lib.TestFlowHelper( yara_flows.YaraDumpProcessMemory.__name__, client_mock, pids=pids or [105], size_limit=size_limit, chunk_size=chunk_size, client_id=self.client_id, ignore_grr_process=True, token=self.token) flow_obj = aff4.FACTORY.Open(session_id, flow.GRRFlow) return flow_obj.ResultCollection()
def testScanAndDump(self): client_mock = action_mocks.MultiGetFileClientMock( yara_actions.YaraProcessScan, yara_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles) procs = [p for p in self.procs if p.pid in [102, 103]] with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid))): session_id = flow_test_lib.TestFlowHelper( yara_flows.YaraProcessScan.__name__, client_mock, yara_signature=test_yara_signature, client_id=self.client_id, token=self.token, include_errors_in_results=True, include_misses_in_results=True, dump_process_on_match=True) flow_obj = aff4.FACTORY.Open(session_id) results = list(flow_obj.ResultCollection()) # 1. Scan result match. # 2. Scan result miss. # 3. ProcDump response. # 3. Stat entry for the dumped file. self.assertEqual(len(results), 4) self.assertIsInstance(results[0], rdf_yara.YaraProcessScanMatch) self.assertIsInstance(results[1], rdf_yara.YaraProcessScanMiss) self.assertIsInstance(results[2], rdf_yara.YaraProcessDumpResponse) self.assertIsInstance(results[3], rdf_client_fs.StatEntry) self.assertEqual(len(results[2].dumped_processes), 1) self.assertEqual(results[0].process.pid, results[2].dumped_processes[0].process.pid) self.assertIn(str(results[2].dumped_processes[0].process.pid), results[3].pathspec.path)
def testConfigFileInclusionWindowsPaths(self): one = r""" Config.includes: - 2.yaml Section1.int: 1 """ two = r""" Section1.int: 2 SecondaryFileIncluded: true """ config_path = "C:\\Windows\\System32\\GRR" def MockedWindowsOpen(filename, _=None): basename = ntpath.basename(filename) dirname = ntpath.dirname(filename) # Make sure we only try to open files from this directory. if dirname != config_path: raise IOError("Tried to open wrong file %s" % filename) if basename == "1.yaml": return io.BytesIO(one.encode("utf-8")) if basename == "2.yaml": return io.BytesIO(two.encode("utf-8")) raise IOError("File not found %s" % filename) # TODO(user): this kind of mocking is a questionable practice at best. # We have Windows-specific tests and should use them for this kind of # testing. # # We need to also use the nt path manipulation modules. with utils.MultiStubber((io, "open", MockedWindowsOpen), (os, "path", ntpath)): conf = self._GetNewConf() conf.Initialize(filename=ntpath.join(config_path, "1.yaml")) self.assertEqual(conf["Section1.int"], 2) self.assertEqual(conf["SecondaryFileIncluded"], True)
def setUp(self): super(ApiIntegrationTest, self).setUp() api_auth_manager.InitializeApiAuthManager() self.token.username = "******" try: webauth.WEBAUTH_MANAGER.SetUserName(self.token.username) except AttributeError: # Only the NullWebAuthManager supports SetUserName pass self.CreateUser(self.token.username) self.port = ApiIntegrationTest.server_port self.endpoint = "http://localhost:%s" % self.port self.api = grr_api.InitHttp(api_endpoint=self.endpoint) poll_stubber = utils.MultiStubber( (grr_api_utils, "DEFAULT_POLL_INTERVAL", 0.1), (grr_api_utils, "DEFAULT_POLL_TIMEOUT", 10)) poll_stubber.Start() self.addCleanup(poll_stubber.Stop)
def _RunYaraProcessScan(self, procs, ignore_grr_process=False, include_errors_in_results=False, include_misses_in_results=False, max_results_per_process=0, **kw): client_mock = action_mocks.ActionMock(memory_actions.YaraProcessScan) with utils.MultiStubber( (psutil, "process_iter", lambda: procs), (psutil, "Process", functools.partial(self.process, procs)), (client_utils, "OpenProcessForMemoryAccess", lambda pid: FakeMemoryProcess(pid=pid))): session_id = flow_test_lib.TestFlowHelper( memory.YaraProcessScan.__name__, client_mock, yara_signature=test_yara_signature, client_id=self.client_id, ignore_grr_process=ignore_grr_process, include_errors_in_results=include_errors_in_results, include_misses_in_results=include_misses_in_results, max_results_per_process=max_results_per_process, token=self.token, **kw) res = flow_test_lib.GetFlowResults(self.client_id.Basename(), session_id) matches = [ r for r in res if isinstance(r, rdf_memory.YaraProcessScanMatch) ] errors = [ r for r in res if isinstance(r, rdf_memory.ProcessMemoryError) ] misses = [ r for r in res if isinstance(r, rdf_memory.YaraProcessScanMiss) ] return (matches, errors, misses)
def testConfigFileInclusionWindowsPaths(self): one = r""" Config.includes: - 2.yaml Section1.int: 1 """ two = r""" Section1.int: 2 SecondaryFileIncluded: true """ config_path = "C:\\Windows\\System32\\GRR" def MockedWindowsOpen(filename, _=None): basename = ntpath.basename(filename) dirname = ntpath.dirname(filename) # Make sure we only try to open files from this directory. if dirname != config_path: raise IOError("Tried to open wrong file %s" % filename) if basename == "1.yaml": # TODO(hanuszczak): YAML, consider using `StringIO`. return io.BytesIO(one) if basename == "2.yaml": # TODO(hanuszczak): YAML, consider using `StringIO`. return io.BytesIO(two) raise IOError("File not found %s" % filename) # We need to also use the nt path manipulation modules. with utils.MultiStubber((__builtin__, "open", MockedWindowsOpen), (os, "path", ntpath)): conf = self._GetNewConf() conf.Initialize(filename=ntpath.join(config_path, "1.yaml")) self.assertEqual(conf["Section1.int"], 2) self.assertEqual(conf["SecondaryFileIncluded"], True)
def Start(self): """Install the stubs.""" modules = { "_winreg": mock.MagicMock(), "ctypes": mock.MagicMock(), "ctypes.wintypes": mock.MagicMock(), # Requires mocking because exceptions.WindowsError does not exist "exceptions": mock.MagicMock(), } self.module_patcher = mock.patch.dict("sys.modules", modules) self.module_patcher.start() # pylint: disable= g-import-not-at-top from grr_response_client.vfs_handlers import registry import exceptions import _winreg # pylint: enable=g-import-not-at-top fixture = RegistryFake() self.stubber = utils.MultiStubber( (registry, "KeyHandle", RegistryFake.FakeKeyHandle), (registry, "OpenKey", fixture.OpenKey), (registry, "QueryValueEx", fixture.QueryValueEx), (registry, "QueryInfoKey", fixture.QueryInfoKey), (registry, "EnumValue", fixture.EnumValue), (registry, "EnumKey", fixture.EnumKey)) self.stubber.Start() # Add the Registry handler to the vfs. vfs.VFSInit().Run() _winreg.HKEY_USERS = "HKEY_USERS" _winreg.HKEY_LOCAL_MACHINE = "HKEY_LOCAL_MACHINE" exceptions.WindowsError = IOError
def _ConfigStub(self, sections=None): mock_config = GetConfigMockClass(sections) return utils.MultiStubber( (config.CONFIG, "GetRaw", mock_config["GetRaw"]), (config.CONFIG, "Get", mock_config["Get"]), (config.CONFIG, "type_infos", mock_config["type_infos"]))
def testXDEV(self): test_dir = os.path.join(self.temp_dir, "xdev_test") local_dev_dir = os.path.join(test_dir, "local_dev") net_dev_dir = os.path.join(test_dir, "net_dev") os.mkdir(test_dir) os.mkdir(local_dev_dir) os.mkdir(net_dev_dir) local_file = os.path.join(local_dev_dir, "local_file") net_file = os.path.join(net_dev_dir, "net_file") with io.open(local_file, "wb") as fd: fd.write(b"local_data") with io.open(net_file, "wb") as fd: fd.write(b"net_data") def MyGetAllowedDevices(xdev, path): if xdev == rdf_file_finder.FileFinderArgs.XDev.ALWAYS: # Never stop at any device boundary. return globbing._XDEV_ALL_ALLOWED elif xdev == rdf_file_finder.FileFinderArgs.XDev.NEVER: base_dev = os.stat(path).st_dev return set([base_dev]) elif xdev == rdf_file_finder.FileFinderArgs.XDev.LOCAL: base_dev = os.stat(path).st_dev # The fake device numbers are base + 5 for local and base + 15 for net. return set([base_dev, base_dev + 5]) else: raise ValueError("Incorrect `xdev` value: %s" % xdev) def MyStat(path): stat_entry = MyStat.old_target(path) if "local_dev" in path or "net_dev" in path: stat_entry_list = list(stat_entry) old_dev = stat_entry.st_dev # Make sure we are modifying the right field. self.assertEqual(old_dev, stat_entry_list[2]) if "local_dev" in path: stat_entry_list[2] = old_dev + 5 else: stat_entry_list[2] = old_dev + 15 return type(stat_entry)(stat_entry_list) return stat_entry with utils.MultiStubber( (os, "stat", MyStat), (globbing, "_GetAllowedDevices", MyGetAllowedDevices)): paths = [test_dir + "/**5"] self.RunAndCheck(paths, expected=[ "local_dev", "local_dev/local_file", "net_dev", "net_dev/net_file" ], unexpected=[], base_path=test_dir, xdev="ALWAYS") self.RunAndCheck( paths, expected=["local_dev", "local_dev/local_file", "net_dev"], unexpected=["net_dev/net_file"], base_path=test_dir, xdev="LOCAL") self.RunAndCheck( paths, expected=["local_dev", "net_dev"], unexpected=["local_dev/local_file", "net_dev/net_file"], base_path=test_dir, xdev="NEVER")
def setUp(self): super(GRRBaseTest, self).setUp() test_user = u"test" system_users_patcher = mock.patch.object( access_control, "SYSTEM_USERS", frozenset(itertools.chain(access_control.SYSTEM_USERS, [test_user]))) system_users_patcher.start() self.addCleanup(system_users_patcher.stop) self.token = access_control.ACLToken(username=test_user, reason="Running tests") self.temp_dir = temp.TempDirPath() config.CONFIG.SetWriteBack( os.path.join(self.temp_dir, "writeback.yaml")) self.addCleanup( lambda: shutil.rmtree(self.temp_dir, ignore_errors=True)) # Each datastore is wrapped with DatabaseValidationWrapper, so we have # to access the delegate directly (assuming it's an InMemoryDB # implementation). data_store.REL_DB.delegate.ClearTestDB() email_alerts.InitializeEmailAlerterOnce() # Stub out the email function self.emails_sent = [] def SendEmailStub(to_user, from_user, subject, message, **unused_kwargs): self.emails_sent.append((to_user, from_user, subject, message)) self.mail_stubber = utils.MultiStubber( (email_alerts.EMAIL_ALERTER, "SendEmail", SendEmailStub), (email.utils, "make_msgid", lambda: "<message id stub>")) self.mail_stubber.Start() self.addCleanup(self.mail_stubber.Stop) # We don't want to send actual email in our tests self.smtp_patcher = mock.patch("smtplib.SMTP") self.mock_smtp = self.smtp_patcher.start() self.addCleanup(self.smtp_patcher.stop) def DisabledSet(*unused_args, **unused_kw): raise NotImplementedError( "Usage of Set() is disabled, please use a configoverrider in tests." ) self.config_set_disable = utils.Stubber(config.CONFIG, "Set", DisabledSet) self.config_set_disable.Start() self.addCleanup(self.config_set_disable.Stop) self._SetupFakeStatsContext() # Turn off WithLimitedCallFrequency-based caching in tests. Tests that need # to test caching behavior explicitly, should turn it on explicitly. with_limited_call_frequency_stubber = utils.Stubber( cache, "WITH_LIMITED_CALL_FREQUENCY_PASS_THROUGH", True) with_limited_call_frequency_stubber.Start() self.addCleanup(with_limited_call_frequency_stubber.Stop)
def setUp(self): super(GRRBaseTest, self).setUp() self.temp_dir = temp.TempDirPath() config.CONFIG.SetWriteBack(os.path.join(self.temp_dir, "writeback.yaml")) self.addCleanup(lambda: shutil.rmtree(self.temp_dir, ignore_errors=True)) data_store.DB.ClearTestDB() # Each datastore is wrapped with DatabaseValidationWrapper, so we have # to access the delegate directly (assuming it's an InMemoryDB # implementation). data_store.REL_DB.delegate.ClearTestDB() aff4.FACTORY.Flush() # Create a Foreman and Filestores, they are used in many tests. aff4_grr.GRRAFF4Init().Run() filestore.FileStoreInit().Run() hunts_results.ResultQueueInitHook().Run() email_alerts.EmailAlerterInit().RunOnce() audit.AuditEventListener._created_logs.clear() # Stub out the email function self.emails_sent = [] def SendEmailStub(to_user, from_user, subject, message, **unused_kwargs): self.emails_sent.append((to_user, from_user, subject, message)) self.mail_stubber = utils.MultiStubber( (email_alerts.EMAIL_ALERTER, "SendEmail", SendEmailStub), (email.utils, "make_msgid", lambda: "<message id stub>")) self.mail_stubber.Start() self.addCleanup(self.mail_stubber.Stop) # We don't want to send actual email in our tests self.smtp_patcher = mock.patch("smtplib.SMTP") self.mock_smtp = self.smtp_patcher.start() self.addCleanup(self.smtp_patcher.stop) def DisabledSet(*unused_args, **unused_kw): raise NotImplementedError( "Usage of Set() is disabled, please use a configoverrider in tests.") self.config_set_disable = utils.Stubber(config.CONFIG, "Set", DisabledSet) self.config_set_disable.Start() self.addCleanup(self.config_set_disable.Stop) if self.use_relational_reads: self.relational_read_stubber = utils.Stubber( data_store, "RelationalDBEnabled", lambda: True) self.relational_read_stubber.Start() self.addCleanup(self.relational_read_stubber.Stop) self._SetupFakeStatsContext() # Turn off WithLimitedCallFrequency-based caching in tests. Tests that need # to test caching behavior explicitly, should turn it on explicitly. with_limited_call_frequency_stubber = utils.Stubber( cache, "WITH_LIMITED_CALL_FREQUENCY_PASS_THROUGH", True) with_limited_call_frequency_stubber.Start() self.addCleanup(with_limited_call_frequency_stubber.Stop)
def MakeRequest(self, instrumentor, manager, path, verify_cb=lambda x: True): with utils.MultiStubber((requests, "request", instrumentor.request), (time, "sleep", instrumentor.sleep)): return manager.OpenServerEndpoint(path, verify_cb=verify_cb)