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)