def testMultiGetFileSetsFileHashAttributeWhenMultipleChunksDownloaded(
self):
client_mock = action_mocks.MultiGetFileClientMock()
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=os.path.join(
self.base_path, "test_img.dd"))
args = transfer.MultiGetFileArgs(pathspecs=[pathspec])
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
h = hashlib.sha256()
with open(os.path.join(self.base_path, "test_img.dd"),
"rb") as model_fd:
h.update(model_fd.read())
cp = db.ClientPath.FromPathSpec(self.client_id, pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.assertEqual(fd_rel_db.hash_id.AsBytes(), h.digest())
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
def testMultiGetFile(self):
"""Test MultiGetFile."""
client_mock = action_mocks.MultiGetFileClientMock()
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=os.path.join(
self.base_path, "test_img.dd"))
args = transfer.MultiGetFileArgs(pathspecs=[pathspec, pathspec])
with test_lib.Instrument(transfer.MultiGetFile,
"StoreStat") as storestat_instrument:
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# We should only have called StoreStat once because the two paths
# requested were identical.
self.assertEqual(len(storestat_instrument.args), 1)
# Fix path for Windows testing.
pathspec.path = pathspec.path.replace("\\", "/")
# Test the AFF4 file that was created.
urn = pathspec.AFF4Path(self.client_id)
fd1 = aff4.FACTORY.Open(urn, token=self.token)
fd2 = open(pathspec.path, "rb")
fd2.seek(0, 2)
self.assertEqual(fd2.tell(), int(fd1.Get(fd1.Schema.SIZE)))
self.CompareFDs(fd1, fd2)
def testMultiGetFileSizeLimit(self):
client_mock = action_mocks.MultiGetFileClientMock()
image_path = os.path.join(self.base_path, "test_img.dd")
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=image_path)
# Read a bit more than one chunk (600 * 1024).
expected_size = 750 * 1024
args = transfer.MultiGetFileArgs(pathspecs=[pathspec],
file_size=expected_size)
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
urn = pathspec.AFF4Path(self.client_id)
blobimage = aff4.FACTORY.Open(urn, token=self.token)
# Make sure a VFSBlobImage got written.
self.assertTrue(isinstance(blobimage, aff4_grr.VFSBlobImage))
self.assertEqual(len(blobimage), expected_size)
data = blobimage.read(100 * expected_size)
self.assertEqual(len(data), expected_size)
expected_data = open(image_path, "rb").read(expected_size)
self.assertEqual(data, expected_data)
hash_obj = data_store_utils.GetFileHashEntry(blobimage)
d = hashlib.sha1()
d.update(expected_data)
expected_hash = d.hexdigest()
self.assertEqual(hash_obj.sha1, expected_hash)
def testMultiGetFileProgressReportsSkippedDuplicatesCorrectly(self):
client_mock = action_mocks.MultiGetFileClientMock()
image_path = os.path.join(self.base_path, "test_img.dd")
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=image_path)
args = transfer.MultiGetFileArgs(pathspecs=[pathspec])
# Let the flow run to make sure the file is collected.
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# Run the flow second time to make sure duplicates are collected.
flow_id = flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
f_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
f_instance = transfer.MultiGetFile(f_obj)
p = f_instance.GetProgress()
self.assertEqual(p.num_collected, 0)
self.assertEqual(p.num_failed, 0)
self.assertEqual(p.num_skipped, 1)
self.assertLen(p.pathspecs_progress, 1)
self.assertEqual(p.pathspecs_progress[0].pathspec, pathspec)
self.assertEqual(p.pathspecs_progress[0].status,
transfer.PathSpecProgress.Status.SKIPPED)
def testExistingChunks(self):
client_mock = action_mocks.MultiGetFileClientMock()
# Make a file to download that is three chunks long.
# For the second run, we change the middle chunk. This will lead to a
# different hash for the whole file and three chunks to download of which we
# already have two.
chunk_size = transfer.MultiGetFile.CHUNK_SIZE
for data in [
"A" * chunk_size + "B" * chunk_size + "C" * 100,
"A" * chunk_size + "X" * chunk_size + "C" * 100
]:
path = os.path.join(self.temp_dir, "test.txt")
with open(path, "wb") as fd:
fd.write(data)
pathspec = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path=path)
args = transfer.MultiGetFileArgs(pathspecs=[pathspec])
flow_test_lib.TestFlowHelper(
transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
urn = pathspec.AFF4Path(self.client_id)
blobimage = aff4.FACTORY.Open(urn)
self.assertEqual(blobimage.size, len(data))
self.assertEqual(blobimage.read(blobimage.size), data)
# Three chunks to get for the first file, only one for the second.
self.assertEqual(client_mock.action_counts["TransferBuffer"], 4)
def testMultiGetFileSetsFileHashAttributeWhenMultipleChunksDownloaded(
self):
client_mock = action_mocks.MultiGetFileClientMock()
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=os.path.join(
self.base_path, "test_img.dd"))
args = transfer.MultiGetFileArgs(pathspecs=[pathspec])
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# Fix path for Windows testing.
pathspec.path = pathspec.path.replace("\\", "/")
# Test the AFF4 file that was created.
urn = pathspec.AFF4Path(self.client_id)
fd_hash = data_store_utils.GetUrnHashEntry(urn)
self.assertTrue(fd_hash)
h = hashlib.sha256()
with open(os.path.join(self.base_path, "test_img.dd"),
"rb") as model_fd:
h.update(model_fd.read())
self.assertEqual(fd_hash.sha256, h.digest())
def testMultiGetFileOfSpecialFiles(self):
"""Test that special /proc/ files are handled correctly.
/proc/ files have the property that they are non seekable from their end
(i.e. seeking them relative to the end is not supported). They also return
an st_size of 0. For example:
$ stat /proc/self/maps
File: '/proc/self/maps'
Size: 0 Blocks: 0 IO Block: 1024 regular empty file
$ head /proc/self/maps
00400000-00409000 r-xp 00000000 fc:01 9180740 /usr/bin/head
00608000-00609000 r--p 00008000 fc:01 9180740 /usr/bin/head
...
When we try to use the MultiGetFile flow, it deduplicates the files and
since it thinks the file has a zero size, the flow will not download the
file, and instead copy the zero size file into it.
"""
client_mock = action_mocks.MultiGetFileClientMock()
# # Create a zero sized file.
zero_sized_filename = os.path.join(self.temp_dir, "zero_size")
with open(zero_sized_filename, "wb") as fd:
pass
pathspec = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path=zero_sized_filename)
flow_test_lib.TestFlowHelper(
transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
file_size="1MiB",
client_id=self.client_id,
pathspecs=[pathspec])
# Now if we try to fetch a real /proc/ filename this will fail because the
# filestore already contains the zero length file
# aff4:/files/nsrl/da39a3ee5e6b4b0d3255bfef95601890afd80709.
pathspec = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path="/proc/self/environ")
flow_test_lib.TestFlowHelper(
transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
file_size=1024 * 1024,
client_id=self.client_id,
pathspecs=[pathspec])
data = open(pathspec.last.path, "rb").read()
# Test the AFF4 file that was created - it should be empty since by default
# we judge the file size based on its stat.st_size.
urn = pathspec.AFF4Path(self.client_id)
fd = aff4.FACTORY.Open(urn, token=self.token)
self.assertEqual(fd.size, len(data))
self.assertMultiLineEqual(fd.read(len(data)), data)
def _RunProcessDump(self, pids=None, size_limit=None):
def FakeProcessMemoryIterator(pid=None): # pylint: disable=invalid-name
del pid
mem_blocks = [
FakeMemoryBlock("A" * 100, 1024),
FakeMemoryBlock("B" * 100, 2048),
]
for m in mem_blocks:
yield m
procs = self.procs
with utils.MultiStubber(
(psutil, "process_iter", lambda: procs),
(psutil, "Process", functools.partial(self.process, procs)),
(yara_procdump, "process_memory_iterator",
FakeProcessMemoryIterator)):
client_mock = action_mocks.MultiGetFileClientMock(
yara_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles)
for s in flow_test_lib.TestFlowHelper(
yara_flows.YaraDumpProcessMemory.__name__,
client_mock,
pids=pids or [105],
size_limit=size_limit,
client_id=self.client_id,
ignore_grr_process=True,
token=self.token):
session_id = s
flow_obj = aff4.FACTORY.Open(session_id, flow.GRRFlow)
return flow_obj.ResultCollection()
def testMultiGetFileDeduplication(self):
client_mock = action_mocks.MultiGetFileClientMock()
pathspecs = []
# Make 10 files to download.
for i in xrange(10):
path = os.path.join(self.temp_dir, "test_%s.txt" % i)
with open(path, "wb") as fd:
fd.write("Hello")
pathspecs.append(
rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=path))
# All those files are the same so the individual chunks should
# only be downloaded once. By forcing maximum_pending_files=1,
# there should only be a single TransferBuffer call.
args = transfer.MultiGetFileArgs(pathspecs=pathspecs,
maximum_pending_files=1)
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
self.assertEqual(client_mock.action_counts["TransferBuffer"], 1)
def testMultiGetFileSizeLimit(self):
client_mock = action_mocks.MultiGetFileClientMock()
image_path = os.path.join(self.base_path, "test_img.dd")
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=image_path)
# Read a bit more than one chunk (600 * 1024).
expected_size = 750 * 1024
args = transfer.MultiGetFileArgs(pathspecs=[pathspec],
file_size=expected_size)
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
expected_data = open(image_path, "rb").read(expected_size)
if data_store.RelationalDBReadEnabled():
cp = db.ClientPath.FromPathSpec(self.client_id.Basename(),
pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.assertEqual(fd_rel_db.size, expected_size)
data = fd_rel_db.read(2 * expected_size)
self.assertLen(data, expected_size)
d = hashlib.sha256()
d.update(expected_data)
self.assertEqual(fd_rel_db.hash_id.AsBytes(), d.digest())
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
else:
urn = pathspec.AFF4Path(self.client_id)
blobimage = aff4.FACTORY.Open(urn, token=self.token)
# Make sure a VFSBlobImage got written.
self.assertIsInstance(blobimage, aff4_grr.VFSBlobImage)
self.assertLen(blobimage, expected_size)
data = blobimage.read(100 * expected_size)
self.assertLen(data, expected_size)
self.assertEqual(data, expected_data)
hash_obj = data_store_utils.GetFileHashEntry(blobimage)
d = hashlib.sha1()
d.update(expected_data)
self.assertEqual(hash_obj.sha1, d.digest())
def testProcessDumpByDefaultErrors(self):
# This tests that not specifying any restrictions on the processes
# to dump does not dump them all which would return tons of data.
client_mock = action_mocks.MultiGetFileClientMock(
memory_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles)
with self.assertRaises(ValueError):
flow_test_lib.TestFlowHelper(memory.DumpProcessMemory.__name__,
client_mock,
client_id=self.client_id,
ignore_grr_process=True,
token=self.token)
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, 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 dumps are not pushed to external file stores.
self.assertEqual(efs.call_count, 0)
results = flow_test_lib.GetFlowResults(self.client_id, 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.assertEmpty(results[2].dumped_processes[0].dump_files)
self.assertLen(results[2].dumped_processes[0].memory_regions, 1)
# TODO: Fix PathSpec.__eq__, then compare PathSpecs here.
self.assertEqual(
results[2].dumped_processes[0].memory_regions[0].file.CollapsePath(
), results[3].pathspec.CollapsePath())
def testProcessDumpByDefaultErrors(self):
# This tests that not specifying any restrictions on the processes
# to dump does not dump them all which would return tons of data.
client_mock = action_mocks.MultiGetFileClientMock(
memory_actions.YaraProcessDump, tempfiles.DeleteGRRTempFiles)
flow_id = flow_test_lib.TestFlowHelper(
memory.DumpProcessMemory.__name__,
client_mock,
client_id=self.client_id,
ignore_grr_process=True,
check_flow_errors=False,
token=self.token)
flow_obj = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(flow_obj.error_message, "No processes to dump specified.")
def testExistingChunks(self):
client_mock = action_mocks.MultiGetFileClientMock()
# Make a file to download that is three chunks long.
# For the second run, we change the middle chunk. This will lead to a
# different hash for the whole file and three chunks to download of which we
# already have two.
chunk_size = transfer.MultiGetFile.CHUNK_SIZE
for data in [
"A" * chunk_size + "B" * chunk_size + "C" * 100,
"A" * chunk_size + "X" * chunk_size + "C" * 100
]:
path = os.path.join(self.temp_dir, "test.txt")
with open(path, "wb") as fd:
fd.write(data)
pathspec = rdf_paths.PathSpec(
pathtype=rdf_paths.PathSpec.PathType.OS, path=path)
args = transfer.MultiGetFileArgs(pathspecs=[pathspec])
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
if data_store.RelationalDBReadEnabled():
cp = db.ClientPath.FromPathSpec(self.client_id.Basename(),
pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.assertEqual(fd_rel_db.size, len(data))
self.assertEqual(fd_rel_db.read(), data)
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
else:
urn = pathspec.AFF4Path(self.client_id)
blobimage = aff4.FACTORY.Open(urn)
self.assertEqual(blobimage.size, len(data))
self.assertEqual(blobimage.read(blobimage.size), data)
# Three chunks to get for the first file, only one for the second.
self.assertEqual(client_mock.action_counts["TransferBuffer"], 4)
def _RunUpdateFlow(self, client_id):
gui_test_lib.CreateFileVersion(
client_id,
"fs/os/c/a.txt",
"Hello World".encode("utf-8"),
timestamp=gui_test_lib.TIME_0)
gui_test_lib.CreateFolder(
client_id, "fs/os/c/TestFolder", timestamp=gui_test_lib.TIME_0)
gui_test_lib.CreateFolder(
client_id, "fs/os/c/bin/TestBinFolder", timestamp=gui_test_lib.TIME_0)
flow_test_lib.FinishAllFlowsOnClient(
client_id,
client_mock=action_mocks.MultiGetFileClientMock(),
check_flow_errors=False)
def testMultiGetFileMultiFiles(self):
"""Test MultiGetFile downloading many files at once."""
client_mock = action_mocks.MultiGetFileClientMock()
pathspecs = []
# Make 30 files to download.
for i in xrange(30):
path = os.path.join(self.temp_dir, "test_%s.txt" % i)
with open(path, "wb") as fd:
fd.write("Hello")
pathspecs.append(
rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=path))
args = transfer.MultiGetFileArgs(pathspecs=pathspecs,
maximum_pending_files=10)
for session_id in flow_test_lib.TestFlowHelper(
transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args):
# Check up on the internal flow state.
flow_obj = aff4.FACTORY.Open(session_id,
mode="r",
token=self.token)
flow_state = flow_obj.state
# All the pathspecs should be in this list.
self.assertEqual(len(flow_state.indexed_pathspecs), 30)
# At any one time, there should not be more than 10 files or hashes
# pending.
self.assertLessEqual(len(flow_state.pending_files), 10)
self.assertLessEqual(len(flow_state.pending_hashes), 10)
# When we finish there should be no pathspecs stored in the flow state.
for flow_pathspec in flow_state.indexed_pathspecs:
self.assertIsNone(flow_pathspec)
for flow_request_data in flow_state.request_data_list:
self.assertIsNone(flow_request_data)
# Now open each file and make sure the data is there.
for pathspec in pathspecs:
urn = pathspec.AFF4Path(self.client_id)
fd = aff4.FACTORY.Open(urn, token=self.token)
self.assertEqual("Hello", fd.read())
def testMultiGetFile(self):
"""Test MultiGetFile."""
client_mock = action_mocks.MultiGetFileClientMock()
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=os.path.join(
self.base_path, "test_img.dd"))
args = transfer.MultiGetFileArgs(pathspecs=[pathspec, pathspec])
with test_lib.Instrument(transfer.MultiGetFileMixin,
"StoreStat") as storestat_instrument:
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# We should only have called StoreStat once because the two paths
# requested were identical.
self.assertLen(storestat_instrument.args, 1)
# Fix path for Windows testing.
pathspec.path = pathspec.path.replace("\\", "/")
fd2 = open(pathspec.path, "rb")
# Test the AFF4 file that was created.
if data_store.RelationalDBReadEnabled():
cp = db.ClientPath.FromPathSpec(self.client_id.Basename(),
pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.CompareFDs(fd2, fd_rel_db)
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
else:
urn = pathspec.AFF4Path(self.client_id)
fd1 = aff4.FACTORY.Open(urn, token=self.token)
fd2.seek(0, 2)
self.assertEqual(fd2.tell(), int(fd1.Get(fd1.Schema.SIZE)))
self.CompareFDs(fd1, fd2)
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, 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)
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(
yara_actions.YaraProcessScan, yara_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(
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)
# 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_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.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 testMultiGetFileSetsFileHashAttributeWhenMultipleChunksDownloaded(
self):
client_mock = action_mocks.MultiGetFileClientMock()
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=os.path.join(
self.base_path, "test_img.dd"))
args = transfer.MultiGetFileArgs(pathspecs=[pathspec])
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
h = hashlib.sha256()
with open(os.path.join(self.base_path, "test_img.dd"),
"rb") as model_fd:
h.update(model_fd.read())
if not data_store.RelationalDBReadEnabled():
# Fix path for Windows testing.
pathspec.path = pathspec.path.replace("\\", "/")
# Test the AFF4 file that was created.
urn = pathspec.AFF4Path(self.client_id)
fd_hash = data_store_utils.GetUrnHashEntry(urn)
self.assertTrue(fd_hash)
self.assertEqual(fd_hash.sha256, h.digest())
if data_store.RelationalDBReadEnabled():
cp = db.ClientPath.FromPathSpec(self.client_id.Basename(),
pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.assertEqual(fd_rel_db.hash_id.AsBytes(), h.digest())
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
def testMultiGetFileDeduplication(self):
client_mock = action_mocks.MultiGetFileClientMock()
pathspecs = []
# Make 10 files to download.
for i in range(10):
path = os.path.join(self.temp_dir, "test_%s.txt" % i)
with open(path, "wb") as fd:
fd.write("Hello")
pathspecs.append(
rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=path))
# All those files are the same so the individual chunks should
# only be downloaded once. By forcing maximum_pending_files=1,
# there should only be a single TransferBuffer call.
args = transfer.MultiGetFileArgs(pathspecs=pathspecs,
maximum_pending_files=1)
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
self.assertEqual(client_mock.action_counts["TransferBuffer"], 1)
if data_store.RelationalDBReadEnabled():
for pathspec in pathspecs:
# Check that each referenced file can be read.
cp = db.ClientPath.FromPathSpec(self.client_id.Basename(),
pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.assertEqual("Hello", fd_rel_db.read())
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
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 _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 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))):
for s in 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):
session_id = s
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.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 _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 testMultiGetFile(self):
"""Test MultiGetFile."""
client_mock = action_mocks.MultiGetFileClientMock()
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=os.path.join(
self.base_path, "test_img.dd"))
expected_size = os.path.getsize(pathspec.path)
args = transfer.MultiGetFileArgs(pathspecs=[pathspec, pathspec])
with test_lib.Instrument(transfer.MultiGetFile,
"_StoreStat") as storestat_instrument:
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# We should only have called StoreStat once because the two paths
# requested were identical.
self.assertLen(storestat_instrument.args, 1)
# Fix path for Windows testing.
pathspec.path = pathspec.path.replace("\\", "/")
with open(pathspec.path, "rb") as fd2:
# Test the file that was created.
cp = db.ClientPath.FromPathSpec(self.client_id, pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.CompareFDs(fd2, fd_rel_db)
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertEqual(history[-1].hash_entry.num_bytes, expected_size)
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
def testMultiGetFileSizeLimit(self):
client_mock = action_mocks.MultiGetFileClientMock()
image_path = os.path.join(self.base_path, "test_img.dd")
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=image_path)
# Read a bit more than one chunk (600 * 1024).
expected_size = 750 * 1024
args = transfer.MultiGetFileArgs(pathspecs=[pathspec],
file_size=expected_size)
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
with open(image_path, "rb") as fd:
expected_data = fd.read(expected_size)
cp = db.ClientPath.FromPathSpec(self.client_id, pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.assertEqual(fd_rel_db.size, expected_size)
data = fd_rel_db.read(2 * expected_size)
self.assertLen(data, expected_size)
d = hashlib.sha256()
d.update(expected_data)
self.assertEqual(fd_rel_db.hash_id.AsBytes(), d.digest())
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertEqual(history[-1].hash_entry.num_bytes, expected_size)
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
def testCreatePerClientFileCollectionHunt(self):
client_ids = self.SetupClients(1)
args = hunt_pb2.ApiCreatePerClientFileCollectionHuntArgs(
description="test hunt")
pca = args.per_client_args.add()
pca.client_id = client_ids[0]
pca.path_type = jobs_pb2.PathSpec.OS
path = os.path.join(self.base_path, "numbers.txt")
pca.paths.append(path)
h = self.api.CreatePerClientFileCollectionHunt(args)
h.Start()
self.RunHunt(client_ids=client_ids,
client_mock=action_mocks.MultiGetFileClientMock())
results = list(h.ListResults())
self.assertLen(results, 1)
self.assertEqual(results[0].client.client_id, client_ids[0])
self.assertEqual(results[0].payload.pathspec.path, path)
self.assertEqual(results[0].payload.pathspec.pathtype,
jobs_pb2.PathSpec.OS)
def testMultiGetFileCorrectlyFetchesSameFileMultipleTimes(self):
"""Test MultiGetFile."""
client_mock = action_mocks.MultiGetFileClientMock()
total_num_chunks = 10
total_size = transfer.MultiGetFile.CHUNK_SIZE * total_num_chunks
path = os.path.join(self.temp_dir, "test_big.txt")
with io.open(path, "wb") as fd:
for i in range(total_num_chunks):
fd.write(
struct.pack("b", i) * transfer.MultiGetFile.CHUNK_SIZE)
pathspec = rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=path)
def _Check(expected_size):
args = transfer.MultiGetFileArgs(pathspecs=[pathspec],
file_size=expected_size)
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# Test the file that was created.
cp = db.ClientPath.FromPathSpec(self.client_id, pathspec)
fd = file_store.OpenFile(cp)
self.assertEqual(fd.size, expected_size)
# Fetch the file twice to test a real-world scenario when a file is first
# fetched with a smaller limit, and then - with a bigger one.
# This tests against a bug in MultiGetFileLogic when first N chunks of
# the file were already fetched during a previous MultiGetFileLogic run,
# and as a consequence the file was considered fully fetched, even if
# the max_file_size value of the current run was much bigger than
# the size of the previously fetched file.
_Check(transfer.MultiGetFileLogic.CHUNK_SIZE * 2)
_Check(total_size)
def testMultiGetFileMultiFiles(self):
"""Test MultiGetFile downloading many files at once."""
client_mock = action_mocks.MultiGetFileClientMock()
pathspecs = []
# Make 30 files to download.
for i in range(30):
path = os.path.join(self.temp_dir, "test_%s.txt" % i)
with io.open(path, "wb") as fd:
fd.write(b"Hello")
pathspecs.append(
rdf_paths.PathSpec(pathtype=rdf_paths.PathSpec.PathType.OS,
path=path))
args = transfer.MultiGetFileArgs(pathspecs=pathspecs,
maximum_pending_files=10)
flow_test_lib.TestFlowHelper(transfer.MultiGetFile.__name__,
client_mock,
token=self.token,
client_id=self.client_id,
args=args)
# Now open each file and make sure the data is there.
for pathspec in pathspecs:
cp = db.ClientPath.FromPathSpec(self.client_id, pathspec)
fd_rel_db = file_store.OpenFile(cp)
self.assertEqual(b"Hello", fd_rel_db.read())
# Check that SHA256 hash of the file matches the contents
# hash and that MD5 and SHA1 are set.
history = data_store.REL_DB.ReadPathInfoHistory(
cp.client_id, cp.path_type, cp.components)
self.assertEqual(history[-1].hash_entry.sha256,
fd_rel_db.hash_id.AsBytes())
self.assertEqual(history[-1].hash_entry.num_bytes, 5)
self.assertIsNotNone(history[-1].hash_entry.sha1)
self.assertIsNotNone(history[-1].hash_entry.md5)
