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")