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)
