def testCallClientChildFlowRace(self):
with test_lib.Instrument(data_store.DB, "MultiSet") as set_function:
session_id = flow.GRRFlow.StartFlow(
client_id=self.client_id,
flow_name="CallClientParentFlow",
token=self.token)
client_request_pos = None
flow_obj_pos = None
for i, (subject, attributes) in enumerate(set_function.args):
for attribute in attributes:
if subject == self.client_id.Queue() and attribute.startswith("task:"):
if client_request_pos is not None:
self.fail("Found multiple request writes, confused.")
client_request_pos = i
if subject == session_id and attribute == "aff4:type":
if flow_obj_pos is not None:
self.fail("Found multiple flow obj writes, confused.")
flow_obj_pos = i
if client_request_pos is None:
self.fail("Client request write not found.")
if flow_obj_pos is None:
self.fail("Flow object write not found.")
# Check that the client request was written after the flow was created.
self.assertLess(flow_obj_pos, client_request_pos,
"The client request was issued before "
"the flow was created.")
def testDiskVolumeInfoWindows(self):
self.client_id = self.SetupClient(0, system="Windows")
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.REGISTRY,
vfs_test_lib.FakeRegistryVFSHandler):
client_mock = action_mocks.WindowsVolumeClientMock()
with test_lib.Instrument(self.flow_base_cls,
"SendReply") as send_reply:
flow_test_lib.TestFlowHelper(
filesystem.DiskVolumeInfo.__name__,
client_mock,
client_id=self.client_id,
token=self.token,
path_list=[r"D:\temp\something", r"/var/tmp"])
results = []
for flow_obj, reply in send_reply.args:
if issubclass(flow_obj.__class__,
filesystem.DiskVolumeInfoMixin):
results.append(reply)
# We asked for D and we guessed systemroot (C) for "/var/tmp", but only
# C and Z are present, so we should just get C.
self.assertCountEqual(
[x.windowsvolume.drive_letter for x in results], ["C:"])
self.assertLen(results, 1)
with test_lib.Instrument(self.flow_base_cls,
"SendReply") as send_reply:
flow_test_lib.TestFlowHelper(
filesystem.DiskVolumeInfo.__name__,
client_mock,
client_id=self.client_id,
token=self.token,
path_list=[r"Z:\blah"])
results = []
for flow_obj, reply in send_reply.args:
if issubclass(flow_obj.__class__,
filesystem.DiskVolumeInfoMixin):
results.append(reply)
self.assertCountEqual(
[x.windowsvolume.drive_letter for x in results], ["Z:"])
self.assertLen(results, 1)
def testProgressThrottling(self):
action = actions.ActionPlugin.classes["ProgressAction"]()
with test_lib.Instrument(client_utils, "KeepAlive") as instrument:
for time, expected_count in [(100, 1), (101, 1), (102, 1), (103, 2),
(104, 2), (105, 2), (106, 3)]:
with test_lib.FakeTime(time):
action.Progress()
self.assertEqual(instrument.call_count, expected_count)
def testFingerprintPresence(self):
client_id = self.SetupClient(0)
path = os.path.join(self.base_path, "winexec_img.dd")
pathspec = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path=path)
pathspec.Append(
path="/winpmem-amd64.sys", pathtype=rdf_paths.PathSpec.PathType.TSK)
client_mock = action_mocks.ActionMock(file_fingerprint.FingerprintFile)
with test_lib.Instrument(flow.GRRFlow, "SendReply") as send_reply:
flow_test_lib.TestFlowHelper(
flows_fingerprint.FingerprintFile.__name__,
client_mock,
token=self.token,
client_id=client_id,
pathspec=pathspec)
self.assertEqual(len(send_reply.args), 1)
for _, reply in send_reply.args:
self.assertTrue(
isinstance(reply, flows_fingerprint.FingerprintFileResult))
self.assertTrue(
str(reply.file_urn).endswith(
"test_data/winexec_img.dd/winpmem-amd64.sys"))
self.assertEqual(
str(reply.hash_entry.sha256),
"40ac571d6d85d669a9a19d498d9f926525481430056ff65746f"
"baf36bee8855f")
self.assertEqual(
str(reply.hash_entry.sha1),
"6e17df1a1020a152f2bf4445d1004b192ae8e42d")
self.assertEqual(
str(reply.hash_entry.md5), "12be1109aa3d3b46c9398972af2008e1")
if data_store.RelationalDBReadEnabled(category="vfs"):
path_info = rdf_objects.PathInfo.FromPathSpec(pathspec)
path_info = data_store.REL_DB.ReadPathInfo(
client_id.Basename(),
path_info.path_type,
components=tuple(path_info.components))
hash_obj = path_info.hash_entry
else:
urn = pathspec.AFF4Path(client_id)
fd = aff4.FACTORY.Open(urn, token=self.token)
self.assertEqual(fd.__class__, aff4_grr.VFSFile)
hash_obj = fd.Get(fd.Schema.HASH)
self.assertEqual(hash_obj.pecoff_sha1,
"1f32fa4eedfba023653c094143d90999f6b9bc4f")
self.assertEqual(hash_obj.signed_data[0].revision, 512)
def testCollectRunKeyBinaries(self):
"""Read Run key from the client_fixtures to test parsing and storage."""
client_id = test_lib.TEST_CLIENT_ID
fixture_test_lib.ClientFixture(client_id, token=self.token)
client = aff4.FACTORY.Open(client_id, token=self.token, mode="rw")
client.Set(client.Schema.SYSTEM("Windows"))
client.Set(client.Schema.OS_VERSION("6.2"))
client_info = client.Get(client.Schema.CLIENT_INFO)
client_info.client_version = config.CONFIG["Source.version_numeric"]
client.Set(client.Schema.CLIENT_INFO, client_info)
client.Flush()
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
vfs_test_lib.FakeFullVFSHandler):
client_mock = action_mocks.ActionMock(
file_fingerprint.FingerprintFile,
searching.Find,
standard.StatFile,
)
# Get KB initialized
for s in flow_test_lib.TestFlowHelper(
artifact.KnowledgeBaseInitializationFlow.__name__,
client_mock,
client_id=client_id,
token=self.token):
session_id = s
col = flow.GRRFlow.ResultCollectionForFID(session_id)
client.Set(client.Schema.KNOWLEDGE_BASE, list(col)[0])
client.Flush()
with test_lib.Instrument(transfer.MultiGetFile,
"Start") as getfile_instrument:
# Run the flow in the emulated way.
for _ in flow_test_lib.TestFlowHelper(
registry.CollectRunKeyBinaries.__name__,
client_mock,
client_id=client_id,
token=self.token):
pass
# Check MultiGetFile got called for our runkey file
download_requested = False
for pathspec in getfile_instrument.args[0][0].args.pathspecs:
if pathspec.path == u"C:\\Windows\\TEMP\\A.exe":
download_requested = True
self.assertTrue(download_requested)
def testCollectRunKeyBinaries(self):
"""Read Run key from the client_fixtures to test parsing and storage."""
client_id = self.SetupClient(0, system="Windows", os_version="6.2")
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
vfs_test_lib.FakeFullVFSHandler):
client_mock = action_mocks.ActionMock(
file_fingerprint.FingerprintFile,
searching.Find,
standard.GetFileStat,
)
# Get KB initialized
session_id = flow_test_lib.TestFlowHelper(
artifact.KnowledgeBaseInitializationFlow.__name__,
client_mock,
client_id=client_id,
token=self.token)
kb = flow_test_lib.GetFlowResults(client_id, session_id)[0]
if data_store.RelationalDBEnabled():
client = data_store.REL_DB.ReadClientSnapshot(
client_id.Basename())
client.knowledge_base = kb
data_store.REL_DB.WriteClientSnapshot(client)
else:
with aff4.FACTORY.Open(client_id, token=self.token,
mode="rw") as client:
client.Set(client.Schema.KNOWLEDGE_BASE, kb)
if data_store.RelationalDBEnabled():
flow_cls = transfer.MultiGetFile
else:
flow_cls = aff4_flows.MultiGetFile
with test_lib.Instrument(flow_cls, "Start") as getfile_instrument:
# Run the flow in the emulated way.
flow_test_lib.TestFlowHelper(
registry.CollectRunKeyBinaries.__name__,
client_mock,
client_id=client_id,
token=self.token)
# Check MultiGetFile got called for our runkey file
download_requested = False
for pathspec in getfile_instrument.args[0][0].args.pathspecs:
if pathspec.path == u"C:\\Windows\\TEMP\\A.exe":
download_requested = True
self.assertTrue(download_requested)
def testProgressThrottling(self):
class MockWorker(object):
def Heartbeat(self):
pass
action = ProgressAction(grr_worker=MockWorker())
with test_lib.Instrument(client_utils, "KeepAlive") as instrument:
for time, expected_count in [(100, 1), (101, 1), (102, 1),
(103, 2), (104, 2), (105, 2),
(106, 3)]:
with test_lib.FakeTime(time):
action.Progress()
self.assertEqual(instrument.call_count, expected_count)
def testBrokenArtifact(self):
"""Tests a broken artifact."""
self.client_id = self.SetupClient(0, system="Windows", os_version="6.2")
client_mock = action_mocks.FileFinderClientMock()
artifact_list = ["BadPathspecArtifact"]
with test_lib.Instrument(transfer.MultiGetFile,
"Start") as getfile_instrument:
flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
client_mock,
artifact_list=artifact_list,
knowledge_base=self._GetKB(),
creator=self.test_username,
client_id=self.client_id,
split_output_by_artifact=True)
self.assertFalse(getfile_instrument.args)