def testGetFilesArchiveGeneratesCorrectArchive(self): client_id, flow_id = self._SetupFlowWithStatEntryResults() blob_size = 1024 * 1024 * 4 blob_data, blob_refs = vfs_test_lib.GenerateBlobRefs(blob_size, "ab") vfs_test_lib.CreateFileWithBlobRefsAndData( db.ClientPath.OS(client_id, ["foo", "bar1"]), blob_refs, blob_data) blob_data, blob_refs = vfs_test_lib.GenerateBlobRefs(blob_size, "cd") vfs_test_lib.CreateFileWithBlobRefsAndData( db.ClientPath.OS(client_id, ["foo", "bar2"]), blob_refs, blob_data) zip_stream = io.BytesIO() self.api.Client(client_id).Flow( flow_id).GetFilesArchive().WriteToStream(zip_stream) zip_fd = zipfile.ZipFile(zip_stream) prefix = "%s_flow_ListProcesses_%s" % (client_id, flow_id) namelist = zip_fd.namelist() self.assertCountEqual(namelist, [ "%s/MANIFEST" % prefix, "%s/%s/client_info.yaml" % (prefix, client_id), "%s/%s/fs/os/foo/bar1" % (prefix, client_id), "%s/%s/fs/os/foo/bar2" % (prefix, client_id), ]) for info in zip_fd.infolist(): self.assertGreater(info.compress_size, 0)
def testGetFilesArchiveDropsStreamingResponsesWhenSecondFileBlobIsMissing( self): client_id, flow_id = self._SetupFlowWithStatEntryResults() blob_data, blob_refs = vfs_test_lib.GenerateBlobRefs( 1024 * 1024 * 4, "abc") vfs_test_lib.CreateFileWithBlobRefsAndData( db.ClientPath.OS(client_id, ["foo", "bar1"]), blob_refs, blob_data[0:2]) zip_stream = io.BytesIO() timestamp = rdfvalue.RDFDatetime.Now() self.api.Client(client_id).Flow( flow_id).GetFilesArchive().WriteToStream(zip_stream) with self.assertRaises(zipfile.BadZipfile): zipfile.ZipFile(zip_stream) # Check that notification was pushed indicating the failure to the user. pending_notifications = list( self.api.GrrUser().ListPendingNotifications( timestamp=timestamp.AsMicrosecondsSinceEpoch())) self.assertLen(pending_notifications, 1) self.assertEqual( pending_notifications[0].data.notification_type, int(rdf_objects.UserNotification.Type. TYPE_FILE_ARCHIVE_GENERATION_FAILED)) self.assertEqual(pending_notifications[0].data.reference.type, pending_notifications[0].data.reference.FLOW) self.assertEqual( pending_notifications[0].data.reference.flow.client_id, client_id) self.assertEqual(pending_notifications[0].data.reference.flow.flow_id, flow_id)
def testGetFilesArchiveDropsStreamingResponsesWhenSecondFileBlobIsMissing( self): blob_data, blob_refs = vfs_test_lib.GenerateBlobRefs(1024 * 1024 * 10, "01") # We write just the references, without actual data, simulating a case # when blobs were not written to the blob store for some reason. vfs_test_lib.CreateFileWithBlobRefsAndData( db.ClientPath.TSK("C.1000000000000000", ["c", "universe", "42"]), blob_refs, blob_data[:1]) zip_stream = io.BytesIO() timestamp = rdfvalue.RDFDatetime.Now() self.api.Client(client_id=self.client_id).File( "fs/tsk/c/universe").GetFilesArchive().WriteToStream(zip_stream) with self.assertRaises(zipfile.BadZipfile): zipfile.ZipFile(zip_stream) # Check that notification was pushed indicating the failure to the user. pending_notifications = list(self.api.GrrUser().ListPendingNotifications( timestamp=timestamp.AsMicrosecondsSinceEpoch())) self.assertLen(pending_notifications, 1) self.assertEqual( pending_notifications[0].data.notification_type, int(rdf_objects.UserNotification.Type .TYPE_FILE_ARCHIVE_GENERATION_FAILED)) self.assertEqual(pending_notifications[0].data.reference.type, pending_notifications[0].data.reference.VFS) self.assertEqual(pending_notifications[0].data.reference.vfs.client_id, self.client_id) self.assertEqual(pending_notifications[0].data.reference.vfs.vfs_path, "fs/tsk/c/universe")
def _WriteFile(self, client_path, blobs_range=None): r_from, r_to = blobs_range or (0, 0) blob_data, blob_refs = vfs_test_lib.GenerateBlobRefs( self.blob_size, "abcdef"[r_from:r_to]) vfs_test_lib.CreateFileWithBlobRefsAndData(client_path, blob_refs, blob_data) return blob_data, blob_refs
def testFilesWithOneBlobAreStillReadToEnsureBlobExists(self): _, long_blob_refs = vfs_test_lib.GenerateBlobRefs(self.blob_size, "cd") _, short_blob_refs1 = vfs_test_lib.GenerateBlobRefs( self.blob_size, "a") _, short_blob_refs2 = vfs_test_lib.GenerateBlobRefs( self.blob_size, "b") path1 = db.ClientPath.OS(self.client_id, ["foo"]) path2 = db.ClientPath.OS(self.client_id, ["bar"]) path3 = db.ClientPath.OS(self.client_id, ["baz"]) # One small file, blob is still read. with mock.patch.object(data_store.BLOBS, "ReadBlobs", wraps=data_store.BLOBS.ReadBlobs) as p: file_store.AddFileWithUnknownHash(path1, short_blob_refs1) p.assert_called_once() # Same for multiple small files. with mock.patch.object(data_store.BLOBS, "ReadBlobs", wraps=data_store.BLOBS.ReadBlobs) as p: file_store.AddFilesWithUnknownHashes({ path1: short_blob_refs1, path2: short_blob_refs2 }) p.assert_called_once() # One large file and two small ones result in a single read for the # all three blobs. with mock.patch.object(data_store.BLOBS, "ReadBlobs", wraps=data_store.BLOBS.ReadBlobs) as p: file_store.AddFilesWithUnknownHashes({ path1: short_blob_refs1, path2: short_blob_refs2, path3: long_blob_refs }) p.assert_called_once() self.assertLen(p.call_args[POSITIONAL_ARGS], 1) self.assertEmpty(p.call_args[KEYWORD_ARGS]) self.assertCountEqual(p.call_args[0][0], [ r.blob_id for r in itertools.chain( short_blob_refs1, short_blob_refs2, long_blob_refs) ])
def setUp(self): super(BlobStreamTest, self).setUp() self.blob_size = 10 self.blob_data, self.blob_refs = vfs_test_lib.GenerateBlobRefs( self.blob_size, "abcde12345") blob_ids = [ref.blob_id for ref in self.blob_refs] data_store.BLOBS.WriteBlobs(dict(zip(blob_ids, self.blob_data))) self.blob_stream = file_store.BlobStream(None, self.blob_refs, None)
def setUp(self): super(AddFileWithUnknownHashTest, self).setUp() self.blob_size = 10 self.blob_data, self.blob_refs = vfs_test_lib.GenerateBlobRefs( self.blob_size, "abcd") blob_ids = [ref.blob_id for ref in self.blob_refs] data_store.BLOBS.WriteBlobs(dict(zip(blob_ids, self.blob_data))) self.client_id = "C.0000111122223333" self.client_path = db.ClientPath.OS(self.client_id, ["foo", "bar"])
def setUp(self): super(OpenFileTest, self).setUp() self.client_id = self.SetupClient(0) self.client_path = db.ClientPath.OS(self.client_id, ("foo", "bar")) blob_size = 10 blob_data, blob_refs = vfs_test_lib.GenerateBlobRefs(blob_size, "abcdef") blob_ids = [ref.blob_id for ref in blob_refs] data_store.BLOBS.WriteBlobs(dict(zip(blob_ids, blob_data))) blob_data, blob_refs = vfs_test_lib.GenerateBlobRefs(blob_size, "def") self.hash_id = file_store.AddFileWithUnknownHash(self.client_path, blob_refs) self.data = b"".join(blob_data) _, blob_refs = vfs_test_lib.GenerateBlobRefs(blob_size, "abc") self.other_hash_id = file_store.AddFileWithUnknownHash( self.client_path, blob_refs) self.invalid_hash_id = rdf_objects.SHA256HashID.FromData(b"")
def testGetFilesArchiveFailsWhenFirstFileBlobIsMissing(self): client_id, flow_id = self._SetupFlowWithStatEntryResults() _, blob_refs = vfs_test_lib.GenerateBlobRefs(10, "0") vfs_test_lib.CreateFileWithBlobRefsAndData( db.ClientPath.OS(client_id, ["foo", "bar1"]), blob_refs, []) zip_stream = io.BytesIO() with self.assertRaisesRegex(grr_api_errors.UnknownError, "Could not find one of referenced blobs"): self.api.Client(client_id).Flow( flow_id).GetFilesArchive().WriteToStream(zip_stream)
def testGetBlobFailsWhenFileIsCorrupt(self): _, blob_refs = vfs_test_lib.GenerateBlobRefs(10, "0") # We write just the references, without actual data, simulating a case # when blobs were not written to the blob store for some reason. vfs_test_lib.CreateFileWithBlobRefsAndData( db.ClientPath.OS("C.1000000000000000", ["c", "bin", "test"]), blob_refs, []) out = io.BytesIO() with self.assertRaises(errors.UnknownError): self.api.Client(client_id=self.client_id).File( "fs/os/c/bin/test").GetBlob().WriteToStream(out)
def testRaisesIfSingleFileChunkIsMissing(self): _, missing_blob_refs = vfs_test_lib.GenerateBlobRefs(self.blob_size, "0") hash_id = rdf_objects.SHA256HashID.FromBytes( missing_blob_refs[0].blob_id.AsBytes()) data_store.REL_DB.WriteHashBlobReferences({hash_id: missing_blob_refs}) client_path = db.ClientPath.OS(self.client_id, ("foo", "bar")) path_info = rdf_objects.PathInfo.OS(components=client_path.components) path_info.hash_entry.sha256 = hash_id.AsBytes() data_store.REL_DB.WritePathInfos(client_path.client_id, [path_info]) # Just getting the generator doesn't raise. chunks = file_store.StreamFilesChunks([client_path]) # Iterating through the generator does actually raise. with self.assertRaises(file_store.BlobNotFoundError): list(chunks)
def testRaisesIfBlobIsMissing(self): _, missing_blob_refs = vfs_test_lib.GenerateBlobRefs( self.blob_size, "0") blob_stream = file_store.BlobStream(None, missing_blob_refs, None) with self.assertRaises(file_store.BlobNotFoundError): blob_stream.read(1)