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)
