def setUp(self):
super(FilterCollectionTest, self).setUp()
self.fd = sequential_collection.GeneralIndexedCollection(
rdfvalue.RDFURN("aff4:/tmp/foo/bar"), token=self.token)
for i in range(10):
self.fd.Add(rdf_paths.PathSpec(path="/var/os/tmp-%d" % i, pathtype="OS"))
def ResultCollectionForArtifact(cls,
session_id,
artifact_name,
token=None):
urn = rdfvalue.RDFURN("_".join(
(str(session_id.Add(flow.RESULTS_SUFFIX)),
utils.SmartStr(artifact_name))))
return sequential_collection.GeneralIndexedCollection(urn, token=token)
def testAddGet(self):
collection = sequential_collection.GeneralIndexedCollection(
rdfvalue.RDFURN("aff4:/sequential_collection/testAddGetIndexed"),
token=self.token)
collection.Add(rdfvalue.RDFInteger(42))
collection.Add(rdfvalue.RDFString("the meaning of life"))
self.assertEqual(collection[0].__class__, rdfvalue.RDFInteger)
self.assertEqual(collection[0], 42)
self.assertEqual(collection[1].__class__, rdfvalue.RDFString)
self.assertEqual(collection[1], "the meaning of life")
def _OpenCollectionPath(coll_path, token=None):
"""Tries to open various types of collections at the given path."""
collection = results.HuntResultCollection(coll_path, token=token)
if collection and collection[0].payload:
return collection
collection = sequential_collection.GeneralIndexedCollection(coll_path,
token=token)
if collection:
return collection
def _OpenCollectionPath(coll_path, token=None):
"""Tries to open various types of collections at the given path."""
coll = aff4.FACTORY.Open(coll_path, token=token)
if coll.__class__.__name__ == "RDFValueCollection":
return coll
collection = results.HuntResultCollection(coll_path, token=token)
if collection and collection[0].payload:
return collection
collection = sequential_collection.GeneralIndexedCollection(coll_path,
token=token)
if collection:
return collection
def testDownloadCollectionIgnoresArtifactResultsWithoutFiles(self):
# Create a collection with URNs to some files.
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
fd.Add(collectors.ArtifactFilesDownloaderResult())
with utils.TempDirectory() as tmpdir:
export_utils.DownloadCollection(self.collection_urn,
tmpdir,
overwrite=True,
dump_client_info=True,
token=self.token,
max_threads=2)
expected_outdir = os.path.join(tmpdir, self.out.Path()[1:])
self.assertFalse(os.path.exists(expected_outdir))
def testDownloadCollectionWithFoldersEntries(self):
"""Check we can download a collection that also references folders."""
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
fd.Add(
rdf_file_finder.FileFinderResult(stat_entry=rdf_client.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile5", pathtype="OS"))))
fd.Add(
rdf_file_finder.FileFinderResult(stat_entry=rdf_client.StatEntry(
pathspec=rdf_paths.PathSpec(path="testdir1", pathtype="OS"),
st_mode=stat.S_IFDIR)))
with utils.TempDirectory() as tmpdir:
export_utils.DownloadCollection(self.collection_urn,
tmpdir,
overwrite=True,
dump_client_info=True,
token=self.token,
max_threads=2)
expected_outdir = os.path.join(tmpdir, self.out.Path()[1:])
# Check we found both files.
self.assertTrue("testfile5" in os.listdir(expected_outdir))
self.assertTrue("testdir1" in os.listdir(expected_outdir))
def testDownloadCollectionWithFlattenOption(self):
"""Check we can download files references in a collection."""
# Create a collection with URNs to some files.
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
fd.Add(rdfvalue.RDFURN(self.out.Add("testfile1")))
fd.Add(
rdf_client.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile2", pathtype="OS")))
fd.Add(
rdf_file_finder.FileFinderResult(stat_entry=rdf_client.StatEntry(
pathspec=rdf_paths.PathSpec(path="testfile5", pathtype="OS"))))
with utils.TempDirectory() as tmpdir:
export_utils.DownloadCollection(self.collection_urn,
tmpdir,
overwrite=True,
dump_client_info=True,
flatten=True,
token=self.token,
max_threads=2)
# Check that "files" folder is filled with symlinks to downloaded files.
symlinks = os.listdir(os.path.join(tmpdir, "files"))
self.assertEqual(len(symlinks), 3)
self.assertListEqual(sorted(symlinks), [
"C.1000000000000000_fs_os_testfile1",
"C.1000000000000000_fs_os_testfile2",
"C.1000000000000000_fs_os_testfile5"
])
self.assertEqual(
os.readlink(
os.path.join(tmpdir, "files",
"C.1000000000000000_fs_os_testfile1")),
os.path.join(tmpdir, "C.1000000000000000", "fs", "os",
"testfile1"))
def testDownloadGeneralIndexedCollection(self):
"""Check we can download files references in GeneralIndexedCollection."""
fd = sequential_collection.GeneralIndexedCollection(
self.collection_urn, token=self.token)
self._AddTestData(fd)
self._VerifyDownload()
