def testGetFlowFilesArchiveReturnsNonLimitedHandlerForArtifactsWhenNeeded( self): router = self._CreateRouter( artifact_collector_flow=rr.RobotRouterArtifactCollectorFlowParams( artifact_collector_flow_name=AnotherArtifactCollector.__name__ ), # pylint: disable=undefined-variable get_flow_files_archive=rr.RobotRouterGetFlowFilesArchiveParams( enabled=True, skip_glob_checks_for_artifact_collector=True, exclude_path_globs=["**/*.txt"], include_only_path_globs=["foo/*", "bar/*"])) flow_id = self._CreateFlowWithRobotId() handler = router.GetFlowFilesArchive( api_flow.ApiGetFlowFilesArchiveArgs(client_id=self.client_id, flow_id=flow_id), context=self.context) self.assertEqual(handler.exclude_path_globs, ["**/*.txt"]) self.assertEqual(handler.include_only_path_globs, ["foo/*", "bar/*"]) flow_id = self._CreateFlowWithRobotId( flow_name=AnotherArtifactCollector.__name__, # pylint: disable=undefined-variable flow_args=rdf_artifacts.ArtifactCollectorFlowArgs( artifact_list=["Foo"])) handler = router.GetFlowFilesArchive( api_flow.ApiGetFlowFilesArchiveArgs(client_id=self.client_id, flow_id=flow_id), context=self.context) self.assertIsNone(handler.exclude_path_globs) self.assertIsNone(handler.include_only_path_globs)
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 testArtifactCollectorIndicatesCollectedSizeAfterCollection(self): registry_stub = artifact_registry.ArtifactRegistry() source = rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.FILE, attributes={ "paths": [os.path.join(self.base_path, "numbers.txt")], }) artifact = rdf_artifacts.Artifact( name="FakeArtifact", sources=[source], doc="fake artifact doc") registry_stub.RegisterArtifact(artifact) client_ref = self.api.Client(client_id=self.client_id) with mock.patch.object(artifact_registry, "REGISTRY", registry_stub): args = rdf_artifacts.ArtifactCollectorFlowArgs( artifact_list=["FakeArtifact"]).AsPrimitiveProto() client_ref.CreateFlow( name=collectors.ArtifactCollectorFlow.__name__, args=args) client_mock = action_mocks.FileFinderClientMock() flow_test_lib.FinishAllFlowsOnClient( self.client_id, client_mock=client_mock) f = client_ref.File("fs/os" + os.path.join(self.base_path, "numbers.txt")).Get() self.assertNotEqual(f.data.hash.sha256, b"") self.assertGreater(f.data.hash.num_bytes, 0) self.assertGreater(f.data.last_collected, 0) self.assertGreater(f.data.last_collected_size, 0)
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 testGetFlowFilesArchiveReturnsNonLimitedHandlerForArtifactsWhenNeeded( self): router = self._CreateRouter( artifact_collector_flow=rr.RobotRouterArtifactCollectorFlowParams( artifact_collector_flow_name=AnotherArtifactCollector.__name__ ), get_flow_files_archive=rr.RobotRouterGetFlowFilesArchiveParams( enabled=True, skip_glob_checks_for_artifact_collector=True, path_globs_blacklist=["**/*.txt"], path_globs_whitelist=["foo/*", "bar/*"])) flow_id = self._CreateFlowWithRobotId() handler = router.GetFlowFilesArchive( api_flow.ApiGetFlowFilesArchiveArgs(client_id=self.client_id, flow_id=flow_id), token=self.token) self.assertEqual(handler.path_globs_blacklist, ["**/*.txt"]) self.assertEqual(handler.path_globs_whitelist, ["foo/*", "bar/*"]) flow_id = self._CreateFlowWithRobotId( flow_name=AnotherArtifactCollector.__name__, flow_args=rdf_artifacts.ArtifactCollectorFlowArgs( artifact_list=["Foo"])) handler = router.GetFlowFilesArchive( api_flow.ApiGetFlowFilesArchiveArgs(client_id=self.client_id, flow_id=flow_id), token=self.token) self.assertTrue(handler.path_globs_blacklist is None) self.assertTrue(handler.path_globs_whitelist is None)
def ArtifactCollectorArgs(self, artifact_list, collect_knowledge_base=False): flow_args = rdf_artifacts.ArtifactCollectorFlowArgs( artifact_list=artifact_list, recollect_knowledge_base=collect_knowledge_base) return collectors.GetArtifactCollectorArgs(flow_args, self.knowledge_base)
def Check(artifacts): router.CreateFlow(api_flow.ApiCreateFlowArgs( flow=api_flow.ApiFlow( name=collectors.ArtifactCollectorFlow.__name__, args=rdf_artifacts.ArtifactCollectorFlowArgs( artifact_list=artifacts)), client_id=self.client_id), token=self.token)
def Check(artifacts): with self.assertRaises(access_control.UnauthorizedAccess): router.CreateFlow(api_flow.ApiCreateFlowArgs( flow=api_flow.ApiFlow( name=collectors.ArtifactCollectorFlow.__name__, args=rdf_artifacts.ArtifactCollectorFlowArgs( artifact_list=artifacts)), client_id=self.client_id), token=self.token)
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 testGrepRegexCombination(self): args = rdf_artifacts.ArtifactCollectorFlowArgs() collect_flow = collectors.ArtifactCollectorFlow( rdf_flow_objects.Flow(args=args)) self.assertEqual(collect_flow._CombineRegex([b"simple"]), b"simple") self.assertEqual(collect_flow._CombineRegex([b"a", b"b"]), b"(a)|(b)") self.assertEqual( collect_flow._CombineRegex([b"a", b"b", b"c"]), b"(a)|(b)|(c)") self.assertEqual( collect_flow._CombineRegex([b"a|b", b"[^_]b", b"c|d"]), b"(a|b)|([^_]b)|(c|d)")
def testInterpolateArgs(self): args = rdf_artifacts.ArtifactCollectorFlowArgs() collect_flow = collectors.ArtifactCollectorFlow( rdf_flow_objects.Flow(args=args)) kb = rdf_client.KnowledgeBase() kb.MergeOrAddUser(rdf_client.User(username="******")) kb.MergeOrAddUser(rdf_client.User(username="******")) collect_flow.state["knowledge_base"] = kb collect_flow.current_artifact_name = "blah" test_rdf = rdf_client.KnowledgeBase() action_args = { "usernames": ["%%users.username%%", "%%users.username%%"], "nointerp": "asdfsdf", "notastring": test_rdf } kwargs = collect_flow.InterpolateDict(action_args) self.assertCountEqual(kwargs["usernames"], ["test1", "test2", "test1", "test2"]) self.assertEqual(kwargs["nointerp"], "asdfsdf") self.assertEqual(kwargs["notastring"], test_rdf) # We should be using an array since users.username will expand to multiple # values. self.assertRaises(ValueError, collect_flow.InterpolateDict, {"bad": "%%users.username%%"}) list_args = collect_flow.InterpolateList( ["%%users.username%%", r"%%users.username%%\aa"]) self.assertCountEqual(list_args, ["test1", "test2", r"test1\aa", r"test2\aa"]) list_args = collect_flow.InterpolateList(["one"]) self.assertEqual(list_args, ["one"]) # Ignore the failure in users.desktop, report the others. collect_flow.args.ignore_interpolation_errors = True list_args = collect_flow.InterpolateList( ["%%users.desktop%%", r"%%users.username%%\aa"]) self.assertCountEqual(list_args, [r"test1\aa", r"test2\aa"]) # Both fail. list_args = collect_flow.InterpolateList( [r"%%users.desktop%%\aa", r"%%users.sid%%\aa"]) self.assertCountEqual(list_args, [])
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 testNotAppliedParsers(self, db: abstract_db.Database) -> None: client_id = db_test_utils.InitializeClient(db) flow_id = "4815162342ABCDEF" flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = collectors.ArtifactCollectorFlow.__name__ flow_obj.args = rdf_artifacts.ArtifactCollectorFlowArgs( apply_parsers=False) db.WriteFlowObject(flow_obj) flow_result = rdf_flow_objects.FlowResult() flow_result.client_id = client_id flow_result.flow_id = flow_id flow_result.tag = "artifact:Fake" flow_result.payload = rdfvalue.RDFString("foobar") db.WriteFlowResults([flow_result]) args = flow_plugin.ApiListFlowApplicableParsersArgs() args.client_id = client_id args.flow_id = flow_id result = self.handler.Handle(args) self.assertCountEqual(result.parsers, [ flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.SINGLE_RESPONSE, name="FakeSingleResponse", ), flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.MULTI_RESPONSE, name="FakeMultiResponse", ), flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.SINGLE_FILE, name="FakeSingleFile", ), flow_plugin.ApiParserDescriptor( type=flow_plugin.ApiParserDescriptor.Type.MULTI_FILE, name="FakeMultiFile", ), ])
def testListFlowApplicableParsers(self): client_id = self.SetupClient(0) flow_id = "4815162342ABCDEF" flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = flow_id flow.flow_class_name = collectors.ArtifactCollectorFlow.__name__ flow.args = rdf_artifacts.ArtifactCollectorFlowArgs( apply_parsers=False) data_store.REL_DB.WriteFlowObject(flow) result = rdf_flow_objects.FlowResult() result.client_id = client_id result.flow_id = flow_id result.tag = "artifact:Fake" result.payload = rdf_client_action.ExecuteResponse(stderr=b"foobar") data_store.REL_DB.WriteFlowResults([result]) class FakeParser(parser.SingleResponseParser[None]): supported_artifacts = ["Fake"] def ParseResponse( self, knowledge_base: rdf_client.KnowledgeBase, response: rdfvalue.RDFValue, ) -> Iterable[None]: raise NotImplementedError() with parser_test_lib._ParserContext("Fake", FakeParser): results = self.api.Client(client_id).Flow( flow_id).ListApplicableParsers() self.assertLen(results.parsers, 1) result = results.parsers[0] self.assertEqual(result.name, "Fake") self.assertEqual(result.type, flow_pb2.ApiParserDescriptor.SINGLE_RESPONSE)
def testAlreadyAppliedParsers(self, db: abstract_db.Database) -> None: client_id = db_test_utils.InitializeClient(db) flow_id = "4815162342ABCDEF" flow_obj = rdf_flow_objects.Flow() flow_obj.client_id = client_id flow_obj.flow_id = flow_id flow_obj.flow_class_name = collectors.ArtifactCollectorFlow.__name__ flow_obj.args = rdf_artifacts.ArtifactCollectorFlowArgs( apply_parsers=True) db.WriteFlowObject(flow_obj) flow_result = rdf_flow_objects.FlowResult() flow_result.client_id = client_id flow_result.flow_id = flow_id flow_result.tag = "artifact:Fake" db.WriteFlowResults([flow_result]) args = flow_plugin.ApiListFlowApplicableParsersArgs() args.client_id = client_id args.flow_id = flow_id result = self.handler.Handle(args) self.assertEmpty(result.parsers)
def testListParsedFlowResults(self): client_id = self.SetupClient(0) flow_id = "4815162342ABCDEF" flow = rdf_flow_objects.Flow() flow.client_id = client_id flow.flow_id = flow_id flow.flow_class_name = collectors.ArtifactCollectorFlow.__name__ flow.args = rdf_artifacts.ArtifactCollectorFlowArgs( apply_parsers=False) flow.persistent_data = {"knowledge_base": rdf_client.KnowledgeBase()} data_store.REL_DB.WriteFlowObject(flow) result = rdf_flow_objects.FlowResult() result.client_id = client_id result.flow_id = flow_id result.tag = "artifact:Echo" response = rdf_client_action.ExecuteResponse() response.stderr = "Lorem ipsum.".encode("utf-8") result.payload = response data_store.REL_DB.WriteFlowResults([result]) response = rdf_client_action.ExecuteResponse() response.stderr = "Dolor sit amet.".encode("utf-8") result.payload = response data_store.REL_DB.WriteFlowResults([result]) class StderrToStdoutParser( parser.SingleResponseParser[rdf_client_action.ExecuteResponse] ): supported_artifacts = ["Echo"] def ParseResponse( self, knowledge_base: rdf_client.KnowledgeBase, response: rdf_client_action.ExecuteResponse, ) -> Iterable[rdf_client_action.ExecuteResponse]: del knowledge_base # Unused. if not isinstance(response, rdf_client_action.ExecuteResponse): raise TypeError( f"Unexpected response type: {type(response)}") parsed_response = rdf_client_action.ExecuteResponse() parsed_response.stdout = response.stderr return [parsed_response] with parser_test_lib._ParserContext("StderrToStdout", StderrToStdoutParser): results = self.api.Client(client_id).Flow( flow_id).ListParsedResults() stdouts = [result.payload.stdout.decode("utf-8") for result in results] self.assertLen(stdouts, 2) self.assertEqual(stdouts[0], "Lorem ipsum.") self.assertEqual(stdouts[1], "Dolor sit amet.")
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")
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")