def testValidateSyntaxInvalidLabel(self): artifact = rdf_artifacts.Artifact(name="Norf", doc="This is Norf.", provides=["domain"], labels=["Mail", "Browser", "Reddit"], supported_os=["Darwin"]) with self.assertRaisesRegexp(rdf_artifacts.ArtifactSyntaxError, "'Reddit'"): ar.ValidateSyntax(artifact)
def testDatabaseArtifactsAreLoadedEvenIfNoDatastoreIsRegistered(self): rel_db = data_store.REL_DB artifact = rdf_artifacts.Artifact(name="Foo") rel_db.WriteArtifact(artifact) registry = ar.ArtifactRegistry() registry.ReloadDatastoreArtifacts() self.assertIsNotNone(registry.GetArtifact("Foo"))
def testValidateSyntaxBrokenProvides(self): artifact = rdf_artifacts.Artifact( name="Thud", doc="This is Thud.", provides=["fqdn", "garbage"], labels=["Network"]) with self.assertRaisesRegexp(rdf_artifacts.ArtifactSyntaxError, "'garbage'"): ar.ValidateSyntax(artifact)
def testValidateSyntaxInvalidSupportedOs(self): artifact = rdf_artifacts.Artifact(name="Quux", doc="This is Quux.", provides=["os"], labels=["Cloud", "Logs"], supported_os=["Solaris"]) with self.assertRaisesRegexp(rdf_artifacts.ArtifactSyntaxError, "'Solaris'"): ar.ValidateSyntax(artifact)
def testOldClientSnapshotFallbackUsesLatestApplicable(self): rel_db = data_store.REL_DB client_id = "C.0123456789abcdef" rel_db.WriteClientMetadata(client_id, first_seen=rdfvalue.RDFDatetime.Now()) # Write some fake snapshot history. kb_0 = rdf_client.KnowledgeBase(os="Linux", os_release="rel0") snapshot_0 = rdf_objects.ClientSnapshot( client_id=client_id, knowledge_base=kb_0) rel_db.WriteClientSnapshot(snapshot_0) kb_1 = rdf_client.KnowledgeBase(os="Linux", os_release="rel1") snapshot_1 = rdf_objects.ClientSnapshot( client_id=client_id, knowledge_base=kb_1) rel_db.WriteClientSnapshot(snapshot_1) kb_2 = rdf_client.KnowledgeBase(os="Linux") snapshot_2 = rdf_objects.ClientSnapshot( client_id=client_id, knowledge_base=kb_2) rel_db.WriteClientSnapshot(snapshot_2) with temp.AutoTempDirPath(remove_non_empty=True) as dirpath: filesystem_test_lib.CreateFile(os.path.join(dirpath, "rel0", "quux")) filesystem_test_lib.CreateFile(os.path.join(dirpath, "rel1", "norf")) # Write a fake artifact. art = rdf_artifacts.Artifact( name="Quux", doc="Lorem ipsum.", sources=[ rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.DIRECTORY, attributes={ "paths": [os.path.join(dirpath, "%%os_release%%", "*")], }), ]) rel_db.WriteArtifact(art) artifact_registry.REGISTRY.ReloadDatastoreArtifacts() flow_id = flow_test_lib.TestFlowHelper( compatibility.GetName(collectors.ArtifactCollectorFlow), client_mock=action_mocks.GlobClientMock(), client_id=client_id, artifact_list=["Quux"], old_client_snapshot_fallback=True, token=self.token) results = flow_test_lib.GetFlowResults(client_id=client_id, flow_id=flow_id) self.assertNotEmpty(results) basenames = [os.path.basename(result.pathspec.path) for result in results] self.assertNotIn("quux", basenames) self.assertIn("norf", basenames)
def testReadArtifactReadsWritten(self): artifact = rdf_artifacts.Artifact(name="Foo", doc="Lorem ipsum dolor sit amet.", urls=["http://example.com/foo"]) self.db.WriteArtifact(artifact) result = self.db.ReadArtifact("Foo") self.assertEqual(result.name, "Foo") self.assertEqual(result.doc, "Lorem ipsum dolor sit amet.") self.assertEqual(result.urls, ["http://example.com/foo"])
def testListArtifacts(self): artifact = rdf_artifacts.Artifact(name='FakeArtifact') registry_stub = artifact_registry.ArtifactRegistry() registry_stub.RegisterArtifact(artifact) data_store.REL_DB.WriteArtifact(artifact) with mock.patch.object(artifact_registry, 'REGISTRY', registry_stub): results = grr_colab.list_artifacts() self.assertLen(results, 1) self.assertEqual(results[0].artifact.name, 'FakeArtifact')
def testGetArtifactPathDependencies(self): sources = [{ "type": rdf_artifacts.ArtifactSource.SourceType.REGISTRY_KEY, "attributes": { "keys": [ r"%%current_control_set%%\Control\Session " r"Manager\Environment\Path" ] } }, { "type": rdf_artifacts.ArtifactSource.SourceType.WMI, "attributes": { "query": "SELECT * FROM Win32_UserProfile " "WHERE SID='%%users.sid%%'" } }, { "type": rdf_artifacts.ArtifactSource.SourceType.GREP, "attributes": { "content_regex_list": ["^%%users.username%%:"] } }] artifact = rdf_artifacts.Artifact( name="artifact", doc="Doco", provides=["environ_windir"], supported_os=["Windows"], urls=["http://blah"], sources=sources) self.assertCountEqual( [x["type"] for x in artifact.ToPrimitiveDict()["sources"]], ["REGISTRY_KEY", "WMI", "GREP"]) class Parser1(object): knowledgebase_dependencies = ["appdata", "sid"] class Parser2(object): knowledgebase_dependencies = ["sid", "desktop"] @classmethod def MockGetClassesByArtifact(unused_cls, _): return [Parser1, Parser2] with utils.Stubber(parser.Parser, "GetClassesByArtifact", MockGetClassesByArtifact): self.assertCountEqual( ar.GetArtifactPathDependencies(artifact), [ "appdata", "sid", "desktop", "current_control_set", "users.sid", "users.username" ])
def testFailuresAreLogged(self): client_id = "C.4815162342abcdef" now = rdfvalue.RDFDatetime.Now() data_store.REL_DB.WriteClientMetadata(client_id=client_id, last_ping=now) snapshot = rdf_objects.ClientSnapshot(client_id=client_id) snapshot.knowledge_base.os = "fakeos" data_store.REL_DB.WriteClientSnapshot(snapshot) raising_artifact_source = rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.COMMAND, attributes={ "cmd": "/bin/echo", "args": ["1"], }) raising_artifact = rdf_artifacts.Artifact( name="RaisingArtifact", doc="Lorem ipsum.", sources=[raising_artifact_source]) registry = artifact_registry.ArtifactRegistry() with mock.patch.object(artifact_registry, "REGISTRY", registry): registry.RegisterArtifact(raising_artifact) flow_id = flow_test_lib.TestFlowHelper( collectors.ArtifactCollectorFlow.__name__, client_mock=action_mocks.ActionMock(standard.ExecuteCommand), client_id=client_id, artifact_list=["RaisingArtifact"], apply_parsers=True, check_flow_errors=True, token=self.token) results = flow_test_lib.GetFlowResults(client_id=client_id, flow_id=flow_id) self.assertLen(results, 1) self.assertEqual(results[0].stdout, "1\n".encode("utf-8")) logs = data_store.REL_DB.ReadFlowLogEntries(client_id=client_id, flow_id=flow_id, offset=0, count=1024) # Log should contain two entries. First one about successful execution of # the command (not interesting), the other one containing the error about # unsuccessful parsing. self.assertLen(logs, 2) self.assertIn("It was bound to happen.", logs[1].message)
def testValidateSyntaxBadSource(self): source = { "type": rdf_artifacts.ArtifactSource.SourceType.ARTIFACT_GROUP, "attributes": {} } artifact = rdf_artifacts.Artifact(name="Barf", doc="This is Barf.", provides=["os"], sources=[source]) with self.assertRaisesRegex(rdf_artifacts.ArtifactSyntaxError, "required attributes"): ar.ValidateSyntax(artifact)
def testDeleteArtifactDeletesMultiple(self): for i in range(42): self.db.WriteArtifact(rdf_artifacts.Artifact(name="Art%s" % i)) for i in range(42): if i % 2 == 0: continue self.db.DeleteArtifact("Art%s" % i) for i in range(42): if i % 2 == 0: self.assertEqual(self.db.ReadArtifact("Art%s" % i).name, "Art%s" % i) else: with self.assertRaises(db.UnknownArtifactError): self.db.ReadArtifact("Art%s" % i)
def testValidateSyntaxBadPathDependency(self): sources = [{ "type": rdf_artifacts.ArtifactSource.SourceType.FILE, "attributes": { "paths": [r"%%systemdrive%%\Temp"] } }] artifact = rdf_artifacts.Artifact(name="bad", doc="Doco", provides=["environ_windir"], supported_os=["Windows"], urls=["http://blah"], sources=sources) with self.assertRaises(rdf_artifacts.ArtifactDefinitionError): ar.ValidateSyntax(artifact)
def testWriteArtifactWritesCopy(self): artifact = rdf_artifacts.Artifact() artifact.name = "Foo" artifact.doc = "Lorem ipsum." self.db.WriteArtifact(artifact) artifact.name = "Bar" artifact.doc = "Dolor sit amet." self.db.WriteArtifact(artifact) foo = self.db.ReadArtifact("Foo") self.assertEqual(foo.name, "Foo") self.assertEqual(foo.doc, "Lorem ipsum.") bar = self.db.ReadArtifact("Bar") self.assertEqual(bar.name, "Bar") self.assertEqual(bar.doc, "Dolor sit amet.")
def testLinuxMountCmdArtifact(self, registry): """Test that LinuxMountCmd artifact can be collected.""" artifact_list = ["LinuxMountCmd"] source = rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.COMMAND, attributes={ "cmd": "/bin/mount", "args": [] }) art_obj = rdf_artifacts.Artifact( name="LinuxMountCmd", doc="Linux output of mount", sources=[source], labels=["System"], supported_os=["Linux"]) art_obj.sources.append(source) registry.RegisterArtifact(art_obj) self.assertTrue(artifact_registry.REGISTRY.GetArtifact("LinuxMountCmd")) # Run the ArtifactCollector to get the expected result. expected = self._RunFlow( collectors.ArtifactCollectorFlow, standard.ExecuteCommand, artifact_list, apply_parsers=False) # The artifact collector receives the same result twice here. # self.assertEqual(len(expected), 1) expected = expected[0] self.assertIsInstance(expected, rdf_client_action.ExecuteResponse) # Run the ClientArtifactCollector to get the actual result. result = self._RunFlow( collectors.ClientArtifactCollector, artifact_collector.ArtifactCollector, artifact_list, apply_parsers=False)[0] self.assertEqual(len(result.collected_artifacts), 1) result = result.collected_artifacts[0].action_results[0].value self.assertIsInstance(result, rdf_client_action.ExecuteResponse) self.assertEqual(result, expected)
def testGetArtifactPathDependencies(self): sources = [{ "type": rdf_artifacts.ArtifactSource.SourceType.REGISTRY_KEY, "attributes": { "keys": [ r"%%current_control_set%%\Control\Session " r"Manager\Environment\Path" ] } }, { "type": rdf_artifacts.ArtifactSource.SourceType.WMI, "attributes": { "query": "SELECT * FROM Win32_UserProfile " "WHERE SID='%%users.sid%%'" } }, { "type": rdf_artifacts.ArtifactSource.SourceType.GREP, "attributes": { "content_regex_list": ["^%%users.username%%:"] } }] artifact = rdf_artifacts.Artifact( name="artifact", doc="Doco", provides=["environ_windir"], supported_os=["Windows"], urls=["http://blah"], sources=sources) self.assertCountEqual( [x["type"] for x in artifact.ToPrimitiveDict()["sources"]], ["REGISTRY_KEY", "WMI", "GREP"]) self.assertCountEqual( ar.GetArtifactPathDependencies(artifact), [ "appdata", "sid", "desktop", "current_control_set", "users.sid", "users.username" ])
def testReadArtifactPatchingDeep(self): source = rdf_artifacts.ArtifactSource() source.type = rdf_artifacts.ArtifactSource.SourceType.REGISTRY_VALUE source.attributes = { b"key_value_pairs": [ { b"key": "foo", b"value": "bar", }, { b"key": b"quux", b"value": 1337, }, ], } artifact = rdf_artifacts.Artifact() artifact.name = "foobar" artifact.doc = "Lorem ipsum." artifact.sources = [source] self.db.WriteArtifact(artifact) artifact = self.db.ReadArtifact("foobar") artifact_registry.Validate(artifact) # Should not raise. self.assertLen(artifact.sources, 1) source = artifact.sources[0] self.assertEqual(source.attributes["key_value_pairs"][0]["key"], "foo") self.assertEqual(source.attributes["key_value_pairs"][0]["value"], "bar") self.assertEqual(source.attributes["key_value_pairs"][1]["key"], "quux") self.assertEqual(source.attributes["key_value_pairs"][1]["value"], 1337) # Read again, to ensure that we retrieve what is stored in the database. artifact = self.db.ReadArtifact("foobar") artifact_registry.Validate(artifact) # Should not raise.
def testValidateSyntax(self): sources = [{ "type": rdf_artifacts.ArtifactSource.SourceType.REGISTRY_KEY, "attributes": { "keys": [ r"%%current_control_set%%\Control\Session " r"Manager\Environment\Path" ] } }, { "type": rdf_artifacts.ArtifactSource.SourceType.FILE, "attributes": { "paths": [r"%%environ_systemdrive%%\Temp"] } }] artifact = rdf_artifacts.Artifact(name="good", doc="Doco", provides=["environ_windir"], supported_os=["Windows"], urls=["http://blah"], sources=sources) ar.ValidateSyntax(artifact)
def testWriteArtifactRaisesWithTooLongName(self): name = "a" * (db.MAX_ARTIFACT_NAME_LENGTH + 1) with self.assertRaises(ValueError): self.db.WriteArtifact(rdf_artifacts.Artifact(name=name))
def testDeletesExistingArtifactsInRelDB(self): data_store.REL_DB.WriteArtifact(rdf_artifacts.Artifact(name="old")) self.assertLen(data_store.REL_DB.ReadAllArtifacts(), 1) migration.MigrateArtifacts() self.assertEmpty(data_store.REL_DB.ReadAllArtifacts())
def testWriteAndReadArtifactWithLongName(self): name = "a" * 1024 self.db.WriteArtifact(rdf_artifacts.Artifact(name=name)) self.assertEqual(self.db.ReadArtifact(name).name, name)
def testWriteArtifactThrowsForEmptyName(self): with self.assertRaises(ValueError): self.db.WriteArtifact(rdf_artifacts.Artifact(name=""))
def testReadArtifactReadsCopy(self): self.db.WriteArtifact(rdf_artifacts.Artifact(name="Foo")) self.db.ReadArtifact("Foo").name = "Bar" self.assertEqual(self.db.ReadArtifact("Foo").name, "Foo")
def testDeleteArtifactDeletesSingle(self): self.db.WriteArtifact(rdf_artifacts.Artifact(name="Foo")) self.db.DeleteArtifact("Foo") with self.assertRaises(db.UnknownArtifactError): self.db.ReadArtifact("Foo")
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")
class ApiListParsedFlowResultsHandlerTest(absltest.TestCase): ECHO1337_ARTIFACT_SOURCE = rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.COMMAND, attributes={ "cmd": "/bin/echo", "args": ["1337"], }) ECHO1337_ARTIFACT = rdf_artifacts.Artifact( name="FakeArtifact", doc="Lorem ipsum.", sources=[ECHO1337_ARTIFACT_SOURCE]) class FakeExecuteCommand(action_mocks.ActionMock): def ExecuteCommand( self, args: rdf_client_action.ExecuteRequest, ) -> Iterable[rdf_client_action.ExecuteResponse]: if args.cmd != "/bin/echo": raise ValueError(f"Unsupported command: {args.cmd}") stdout = " ".join(args.args).encode("utf-8") return [rdf_client_action.ExecuteResponse(stdout=stdout)] def setUp(self): super(ApiListParsedFlowResultsHandlerTest, self).setUp() self.handler = flow_plugin.ApiListParsedFlowResultsHandler() @db_test_lib.WithDatabase def testValidatesFlowName(self, db: abstract_db.Database): token = _CreateToken(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=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, "artifact-collector"): self.handler.Handle(args, token=token) @db_test_lib.WithDatabase 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) @db_test_lib.WithDatabase def testParsesArtifactCollectionResults(self, db: abstract_db.Database): token = _CreateToken(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, token=token) flow_test_lib.FinishAllFlowsOnClient(client_id) 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 = 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, 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, b"4815162342") @db_test_lib.WithDatabase def testReportsArtifactCollectionErrors(self, db: abstract_db.Database): token = _CreateToken(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, token=token) flow_test_lib.FinishAllFlowsOnClient(client_id) 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]: 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, token=token) self.assertEmpty(result.items) self.assertLen(result.errors, 1) self.assertEqual(result.errors[0], "Lorem ipsum.") @db_test_lib.WithDatabase 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") @db_test_lib.WithDatabase 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") @db_test_lib.WithDatabase 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 testWriteAndReadArtifactWithUnicodeName(self): name = "🍻foo🍻" self.db.WriteArtifact(rdf_artifacts.Artifact(name=name)) self.assertEqual(self.db.ReadArtifact(name).name, name)
def testWriteAndReadArtifactWithLongName(self): name = "x" + "🧙" * (db.MAX_ARTIFACT_NAME_LENGTH - 2) + "x" self.db.WriteArtifact(rdf_artifacts.Artifact(name=name)) self.assertEqual(self.db.ReadArtifact(name).name, name)
def testWriteArtifactMany(self): for i in range(42): self.db.WriteArtifact(rdf_artifacts.Artifact(name="Art%s" % i)) for i in range(42): self.assertEqual(self.db.ReadArtifact("Art%s" % i).name, "Art%s" % i)