def CheckFilesDownloaded(self, fnames):
for fname in fnames:
if data_store.RelationalDBEnabled():
path_info = data_store.REL_DB.ReadPathInfo(
self.client_id.Basename(),
rdf_objects.PathInfo.PathType.OS,
components=self.FilenameToPathComponents(fname))
size = path_info.stat_entry.st_size
else:
file_urn = self.FileNameToURN(fname)
with aff4.FACTORY.Open(file_urn, token=self.token) as fd:
size = fd.Get(fd.Schema.SIZE)
with io.open(os.path.join(self.base_path, "searching", fname)) as fd:
test_data = fd.read()
self.assertEqual(size, len(test_data))
if data_store.RelationalDBEnabled():
fd = file_store.OpenFile(
db.ClientPath(
self.client_id.Basename(),
rdf_objects.PathInfo.PathType.OS,
components=self.FilenameToPathComponents(fname)))
# Make sure we can actually read the file.
self.assertEqual(fd.read(), test_data)
def Handle(self, args, token=None):
decoder = decoders.FACTORY.Create(args.decoder_name)
path_type, components = rdf_objects.ParseCategorizedPath(args.file_path)
client_path = db.ClientPath(str(args.client_id), path_type, components)
fd = file_store.OpenFile(client_path)
return api_call_handler_base.ApiBinaryStream(
filename=client_path.components[-1],
content_generator=decoder.Decode(fd))
def _VerifyDownloadedFiles(self, results):
if data_store.RelationalDBEnabled():
for r in results:
original_path = r.stat_entry.pathspec.path
fd = file_store.OpenFile(
db.ClientPath(
self.client_id.Basename(),
rdf_objects.PathInfo.PathType.OS,
components=original_path.strip("/").split("/")))
with io.open(original_path, "rb") as orig_fd:
self.assertEqual(fd.read(), orig_fd.read())
else:
for r in results:
file_urn = r.stat_entry.AFF4Path(self.client_id)
fd = aff4.FACTORY.Open(file_urn, token=self.token)
data = fd.read()
with io.open(r.stat_entry.pathspec.path, "rb") as orig_fd:
self.assertEqual(data, orig_fd.read())