def testValidatesParsersWereNotApplied(self, db: abstract_db.Database): token = _CreateToken(db) client_id = db_test_utils.InitializeClient(db) with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(self.ECHO1337_ARTIFACT) flow_args = rdf_artifacts.ArtifactCollectorFlowArgs() flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name] flow_args.apply_parsers = True flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, self.FakeExecuteCommand(), client_id=client_id, args=flow_args, token=token) flow_test_lib.FinishAllFlowsOnClient(client_id) args = flow_plugin.ApiListParsedFlowResultsArgs() args.client_id = client_id args.flow_id = flow_id with self.assertRaisesRegex(ValueError, "already parsed"): self.handler.Handle(args, token=token)
def testEmptyResults(self, db: abstract_db.Database): token = _CreateToken(db) client_id = db_test_utils.InitializeClient(db) fake_artifact = rdf_artifacts.Artifact(name="FakeArtifact", doc="Lorem ipsum.", sources=[]) with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(fake_artifact) flow_args = rdf_artifacts.ArtifactCollectorFlowArgs() flow_args.artifact_list = [fake_artifact.name] flow_args.apply_parsers = False flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, self.FakeExecuteCommand(), client_id=client_id, args=flow_args, token=token) flow_test_lib.FinishAllFlowsOnClient(client_id) 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.assertEmpty(result.items)
def testValidatesFlowName(self, db: abstract_db.Database): context = _CreateContext(db) class FakeFlow(flow_base.FlowBase): def Start(self): self.CallState("End") def End(self, responses: flow_responses.Responses) -> None: del responses # Unused. client_id = db_test_utils.InitializeClient(db) flow_id = flow_test_lib.TestFlowHelper( FakeFlow.__name__, client_id=client_id, token=access_control.ACLToken(username=context.username)) flow_test_lib.FinishAllFlowsOnClient(client_id) args = flow_plugin.ApiListParsedFlowResultsArgs() args.client_id = client_id args.flow_id = flow_id with self.assertRaisesRegex(ValueError, "artifact-collector"): self.handler.Handle(args, context=context)
def testParsesArtifactCollectionResults(self, db: abstract_db.Database): context = _CreateContext(db) with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(self.ECHO1337_ARTIFACT) flow_args = rdf_artifacts.ArtifactCollectorFlowArgs() flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name] flow_args.apply_parsers = False client_id = db_test_utils.InitializeClient(db) flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, self.FakeExecuteCommand(), client_id=client_id, args=flow_args, creator=context.username) flow_test_lib.FinishAllFlowsOnClient(client_id) class FakeParser( abstract_parser.SingleResponseParser[ rdf_client_action.ExecuteResponse], ): supported_artifacts = [self.ECHO1337_ARTIFACT.name] def ParseResponse( self, knowledge_base: rdf_client.KnowledgeBase, response: rdf_client_action.ExecuteResponse, ) -> Iterable[rdf_client_action.ExecuteResponse]: precondition.AssertType(response, rdf_client_action.ExecuteResponse) parsed_response = rdf_client_action.ExecuteResponse() parsed_response.stdout = response.stdout parsed_response.stderr = b"4815162342" return [parsed_response] with parser_test_lib._ParserContext("Fake", FakeParser): args = flow_plugin.ApiListParsedFlowResultsArgs( client_id=client_id, flow_id=flow_id, offset=0, count=1024) result = self.handler.Handle(args, context=context) self.assertEmpty(result.errors) self.assertLen(result.items, 1) response = result.items[0].payload self.assertIsInstance(response, rdf_client_action.ExecuteResponse) self.assertEqual(response.stdout, b"1337") self.assertEqual(response.stderr, b"4815162342")
def testReportsArtifactCollectionErrors(self, db: abstract_db.Database): context = _CreateContext(db) with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(self.ECHO1337_ARTIFACT) flow_args = rdf_artifacts.ArtifactCollectorFlowArgs() flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name] flow_args.apply_parsers = False client_id = db_test_utils.InitializeClient(db) flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, self.FakeExecuteCommand(), client_id=client_id, args=flow_args, creator=context.username) flow_test_lib.FinishAllFlowsOnClient(client_id) class FakeParser( abstract_parser.SingleResponseParser[ rdf_client_action.ExecuteResponse], ): supported_artifacts = [self.ECHO1337_ARTIFACT.name] def ParseResponse( self, knowledge_base: rdf_client.KnowledgeBase, response: rdf_client_action.ExecuteResponse ) -> Iterable[rdf_client_action.ExecuteResponse]: del knowledge_base, response # Unused. raise abstract_parser.ParseError("Lorem ipsum.") with parser_test_lib._ParserContext("Fake", FakeParser): args = flow_plugin.ApiListParsedFlowResultsArgs( client_id=client_id, flow_id=flow_id, offset=0, count=1024) result = self.handler.Handle(args, context=context) self.assertEmpty(result.items) self.assertLen(result.errors, 1) self.assertEqual(result.errors[0], "Lorem ipsum.")
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")
def testUsesKnowledgebaseFromFlow(self, db: abstract_db.Database): token = _CreateToken(db) client_id = db_test_utils.InitializeClient(db) # This is the snapshot that is visible to the flow and should be used for # parsing results. snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id snapshot.knowledge_base.os = "redox" db.WriteClientSnapshot(snapshot) with mock.patch.object( artifact_registry, "REGISTRY", artifact_registry.ArtifactRegistry()) as registry: registry.RegisterArtifact(self.ECHO1337_ARTIFACT) flow_args = rdf_artifacts.ArtifactCollectorFlowArgs() flow_args.artifact_list = [self.ECHO1337_ARTIFACT.name] flow_args.apply_parsers = False flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, self.FakeExecuteCommand(), client_id=client_id, args=flow_args, token=token) class FakeParser(abstract_parser.SingleResponseParser): supported_artifacts = [self.ECHO1337_ARTIFACT.name] def ParseResponse( self, knowledge_base: rdf_client.KnowledgeBase, response: rdf_client_action.ExecuteResponse, ) -> Iterable[rdf_client_action.ExecuteResponse]: precondition.AssertType(response, rdf_client_action.ExecuteResponse) parsed_response = rdf_client_action.ExecuteResponse() parsed_response.stdout = response.stdout parsed_response.stderr = knowledge_base.os.encode("utf-8") return [parsed_response] # This is a snapshot written to the database after the responses were # collected, so this should not be used for parsing. snapshot = rdf_objects.ClientSnapshot() snapshot.client_id = client_id snapshot.knowledge_base.os = "linux" db.WriteClientSnapshot(snapshot) with parser_test_lib._ParserContext("Fake", FakeParser): 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.assertIsInstance(response, rdf_client_action.ExecuteResponse) self.assertEqual(response.stdout, b"1337") self.assertEqual(response.stderr.decode("utf-8"), "redox")