def setUp(self):
super(TestFlowManagement, self).setUp()
self.client_id = self.SetupClient(0, fqdn="Host000011112222")
self.RequestAndGrantClientApproval(self.client_id)
self.action_mock = action_mocks.FileFinderClientMock()
def testFileFinderWorkflowWorks(self):
self.InitRouterConfig(self.__class__.FILE_FINDER_ROUTER_CONFIG %
self.token.username)
client_ref = self.api.Client(client_id=self.client_id.Basename())
args = rdf_file_finder.FileFinderArgs(
paths=[
os.path.join(self.base_path, "test.plist"),
os.path.join(self.base_path, "numbers.txt"),
os.path.join(self.base_path, "numbers.txt.ver2")
],
action=rdf_file_finder.FileFinderAction.Download(
)).AsPrimitiveProto()
flow_obj = client_ref.CreateFlow(name=file_finder.FileFinder.__name__,
args=args)
self.assertEqual(flow_obj.data.state, flow_obj.data.RUNNING)
# Now run the flow we just started.
client_id = rdf_client.ClientURN(flow_obj.client_id)
flow_urn = client_id.Add("flows").Add(flow_obj.flow_id)
flow_test_lib.TestFlowHelper(
flow_urn,
client_id=client_id,
client_mock=action_mocks.FileFinderClientMock(),
token=self.token)
# Refresh flow.
flow_obj = client_ref.Flow(flow_obj.flow_id).Get()
self.assertEqual(flow_obj.data.state, flow_obj.data.TERMINATED)
# Check that we got 3 results (we downloaded 3 files).
results = list(flow_obj.ListResults())
self.assertEqual(len(results), 3)
# We expect results to be FileFinderResult.
self.assertItemsEqual([
os.path.basename(r.payload.stat_entry.pathspec.path)
for r in results
], ["test.plist", "numbers.txt", "numbers.txt.ver2"])
# Now downloads the files archive.
zip_stream = io.BytesIO()
flow_obj.GetFilesArchive().WriteToStream(zip_stream)
zip_fd = zipfile.ZipFile(zip_stream)
# Now check that the archive has only "test.plist" file, as it's the
# only file that matches the whitelist (see FILE_FINDER_ROUTER_CONFIG).
# There should be 3 items in the archive: the hash of the "test.plist"
# file, the symlink to this hash and the MANIFEST file.
namelist = zip_fd.namelist()
self.assertEqual(len(namelist), 3)
# First component of every path in the archive is the containing folder,
# we should strip it.
namelist = [os.path.join(*n.split(os.sep)[1:]) for n in namelist]
self.assertEqual(
sorted([
# pyformat: disable
os.path.join(self.client_id.Basename(), "fs", "os",
self.base_path.strip("/"), "test.plist"),
os.path.join(self.client_id.Basename(), "client_info.yaml"),
"MANIFEST"
# pyformat: enable
]),
sorted(namelist))
def setUp(self): super(TestFlowNotifications, self).setUp() self.client_id = self.SetupClient(0).Basename() self.RequestAndGrantClientApproval(self.client_id) self.action_mock = action_mocks.FileFinderClientMock()
def setUp(self):
super().setUp()
self.client_id = self.SetupClient(0)
self.RequestAndGrantClientApproval(self.client_id)
self.action_mock = action_mocks.FileFinderClientMock()
def testUsesCollectionTimeFiles(self, db: abstract_db.Database):
token = _CreateToken(db)
client_id = db_test_utils.InitializeClient(db)
snapshot = rdf_objects.ClientSnapshot()
snapshot.client_id = client_id
snapshot.knowledge_base.os = "redox"
db.WriteClientSnapshot(snapshot)
with temp.AutoTempFilePath() as temp_filepath:
fake_artifact_source = rdf_artifacts.ArtifactSource(
type=rdf_artifacts.ArtifactSource.SourceType.FILE,
attributes={
"paths": [temp_filepath],
})
fake_artifact = rdf_artifacts.Artifact(
name="FakeArtifact",
doc="Lorem ipsum.",
sources=[fake_artifact_source])
flow_args = rdf_artifacts.ArtifactCollectorFlowArgs()
flow_args.artifact_list = [fake_artifact.name]
flow_args.apply_parsers = False
with io.open(temp_filepath, mode="wb") as temp_filedesc:
temp_filedesc.write(b"OLD")
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(fake_artifact)
# First, we run the artifact collector to collect the old file and save
# the flow id to parse the results later.
flow_id = flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
action_mocks.FileFinderClientMock(),
client_id=client_id,
args=flow_args,
token=token)
flow_test_lib.FinishAllFlowsOnClient(client_id)
with io.open(temp_filepath, mode="wb") as temp_filedesc:
temp_filedesc.write(b"NEW")
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(fake_artifact)
# Now, we run the artifact collector again to collect the new file to
# update to this version on the server. The parsing should be performed
# against the previous flow.
flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
action_mocks.FileFinderClientMock(),
client_id=client_id,
args=flow_args,
token=token)
flow_test_lib.FinishAllFlowsOnClient(client_id)
class FakeFileParser(abstract_parser.SingleFileParser):
supported_artifacts = [fake_artifact.name]
def ParseFile(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspec: rdf_paths.PathSpec,
filedesc: file_store.BlobStream,
) -> Iterable[rdfvalue.RDFBytes]:
del knowledge_base, pathspec # Unused.
return [rdfvalue.RDFBytes(filedesc.Read())]
with parser_test_lib._ParserContext("FakeFile", FakeFileParser):
args = flow_plugin.ApiListParsedFlowResultsArgs(
client_id=client_id, flow_id=flow_id, offset=0, count=1024)
result = self.handler.Handle(args, token=token)
self.assertEmpty(result.errors)
self.assertLen(result.items, 1)
response = result.items[0].payload
self.assertEqual(response, b"OLD")
