def testSingleResponseAndSingleFileParser(self):
class FooParser(parsers.SingleResponseParser[rdfvalue.RDFString]):
supported_artifacts = ["Quux"]
def ParseResponse(
self,
knowledge_base: rdf_client.KnowledgeBase,
response: rdfvalue.RDFValue,
) -> Iterator[rdfvalue.RDFString]:
del knowledge_base # Unused.
if not isinstance(response, rdfvalue.RDFString):
raise TypeError(f"Unexpected response type: {type(response)}")
yield rdfvalue.RDFString(f"FOO-{response}")
class BarParser(parsers.SingleFileParser[rdfvalue.RDFString]):
supported_artifacts = ["Quux"]
def ParseFile(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspec: rdf_paths.PathSpec,
filedesc: IO[bytes],
) -> Iterator[rdfvalue.RDFString]:
del knowledge_base, pathspec, filedesc # Unused.
yield rdfvalue.RDFString("BAR")
with parser_test_lib._ParserContext("Foo", FooParser):
with parser_test_lib._ParserContext("Bar", BarParser):
factory = parsers.ArtifactParserFactory("Quux")
knowledge_base = rdf_client.KnowledgeBase()
applicator = artifact.ParserApplicator(
factory, client_id=self.client_id, knowledge_base=knowledge_base)
applicator.Apply([
rdfvalue.RDFString("THUD"),
rdfvalue.RDFString("BLARGH"),
])
responses = list(applicator.Responses())
self.assertLen(responses, 2)
self.assertEqual(responses[0], rdfvalue.RDFString("FOO-THUD"))
self.assertEqual(responses[1], rdfvalue.RDFString("FOO-BLARGH"))
def testApplySingleResponseError(self):
class FooParseError(parsers.ParseError):
pass
class FooParser(parsers.SingleResponseParser):
supported_artifacts = ["Foo"]
def ParseResponse(
self,
knowledge_base: rdf_client.KnowledgeBase,
response: rdf_client_fs.StatEntry,
) -> Iterable[rdfvalue.RDFString]:
del knowledge_base, response # Unused.
raise FooParseError("Lorem ipsum.")
with parser_test_lib._ParserContext("Foo", FooParser):
factory = parsers.ArtifactParserFactory("Foo")
client_id = self.client_id
knowledge_base = rdf_client.KnowledgeBase()
applicator = artifact.ParserApplicator(factory, client_id,
knowledge_base)
applicator.Apply([rdf_client_fs.StatEntry()])
errors = list(applicator.Errors())
self.assertLen(errors, 1)
self.assertIsInstance(errors[0], FooParseError)
responses = list(applicator.Responses())
self.assertEmpty(responses)
def testApplyMultipleParsersError(self):
class QuuxParseError(parsers.ParseError):
pass
class QuuxParser(parsers.MultiResponseParser):
supported_artifacts = ["Quux"]
def ParseResponses(
self,
knowledge_base: rdf_client.KnowledgeBase,
responses: Collection[rdf_client_fs.StatEntry],
) -> Iterable[rdfvalue.RDFInteger]:
del knowledge_base, responses # Unused.
raise QuuxParseError("Lorem ipsum.")
with parser_test_lib._ParserContext("Quux", QuuxParser):
factory = parsers.ArtifactParserFactory("Quux")
client_id = self.client_id
knowledge_base = rdf_client.KnowledgeBase()
applicator = artifact.ParserApplicator(factory, client_id,
knowledge_base)
applicator.Apply([rdf_client_fs.StatEntry()])
errors = list(applicator.Errors())
self.assertLen(errors, 1)
self.assertIsInstance(errors[0], QuuxParseError)
responses = list(applicator.Responses())
self.assertEmpty(responses)
def testSingleFileResponse(self):
class NorfParser(parsers.SingleFileParser):
supported_artifacts = ["Norf"]
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("Norf", NorfParser):
factory = parsers.ArtifactParserFactory("Norf")
client_id = self.client_id
knowledge_base = rdf_client.KnowledgeBase()
stat_entry = rdf_client_fs.StatEntry()
stat_entry.pathspec.path = "foo/bar/baz"
stat_entry.pathspec.pathtype = rdf_paths.PathSpec.PathType.OS
self._WriteFile(stat_entry.pathspec.path, b"4815162342")
applicator = artifact.ParserApplicator(factory, client_id,
knowledge_base)
applicator.Apply([stat_entry])
errors = list(applicator.Errors())
self.assertEmpty(errors)
responses = list(applicator.Responses())
self.assertLen(responses, 1)
self.assertEqual(responses[0], b"4815162342")
def testApplySingleResponseSuccessful(self):
class FooParser(parsers.SingleResponseParser):
supported_artifacts = ["Foo"]
def ParseResponse(
self,
knowledge_base: rdf_client.KnowledgeBase,
response: rdf_client_fs.StatEntry,
) -> Iterable[rdfvalue.RDFString]:
return [
rdfvalue.RDFString(
f"{knowledge_base.os}:{response.st_dev}")
]
with parser_test_lib._ParserContext("Foo", FooParser):
factory = parsers.ArtifactParserFactory("Foo")
client_id = self.client_id
knowledge_base = rdf_client.KnowledgeBase(os="Redox")
applicator = artifact.ParserApplicator(factory, client_id,
knowledge_base)
applicator.Apply([rdf_client_fs.StatEntry(st_dev=1337)])
errors = list(applicator.Errors())
self.assertEmpty(errors)
responses = list(applicator.Responses())
self.assertEqual(responses, ["Redox:1337"])
def testApplyMultiResponseSuccess(self):
class QuuxParser(parsers.MultiResponseParser):
supported_artifacts = ["Quux"]
def ParseResponses(
self,
knowledge_base: rdf_client.KnowledgeBase,
responses: Collection[rdf_client_fs.StatEntry],
) -> Iterable[rdfvalue.RDFInteger]:
return [stat_entry.st_dev for stat_entry in responses]
with parser_test_lib._ParserContext("Quux", QuuxParser):
factory = parsers.ArtifactParserFactory("Quux")
client_id = self.client_id
knowledge_base = rdf_client.KnowledgeBase()
applicator = artifact.ParserApplicator(factory, client_id,
knowledge_base)
applicator.Apply([
rdf_client_fs.StatEntry(st_dev=42),
rdf_client_fs.StatEntry(st_dev=1337),
])
errors = list(applicator.Errors())
self.assertEmpty(errors)
responses = list(applicator.Responses())
self.assertCountEqual(responses, [42, 1337])
def testSingleResponseAndSingleFileParserWithStatResponse(self):
class FooParser(parsers.SingleResponseParser[rdfvalue.RDFString]):
supported_artifacts = ["Quux"]
def ParseResponse(
self,
knowledge_base: rdf_client.KnowledgeBase,
response: rdfvalue.RDFValue,
) -> Iterator[rdfvalue.RDFString]:
del knowledge_base # Unused.
if not isinstance(response, rdf_client_fs.StatEntry):
raise TypeError(f"Unexpected response type: {type(response)}")
yield rdfvalue.RDFString(f"PATH('{response.pathspec.path}')")
class BarParser(parsers.SingleFileParser[rdfvalue.RDFString]):
supported_artifacts = ["Quux"]
def ParseFile(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspec: rdf_paths.PathSpec,
filedesc: IO[bytes],
) -> Iterator[rdfvalue.RDFString]:
raise NotImplementedError()
with parser_test_lib._ParserContext("Foo", FooParser):
with parser_test_lib._ParserContext("Bar", BarParser):
factory = parsers.ArtifactParserFactory("Quux")
knowledge_base = rdf_client.KnowledgeBase()
stat_entry = rdf_client_fs.StatEntry()
stat_entry.pathspec.path = "foo/bar/baz"
stat_entry.pathspec.pathtype = rdf_paths.PathSpec.PathType.OS
applicator = artifact.ParserApplicator(
factory, client_id=self.client_id, knowledge_base=knowledge_base)
applicator.Apply([stat_entry])
responses = list(applicator.Responses())
self.assertLen(responses, 1)
self.assertEqual(responses[0], "PATH('foo/bar/baz')")
def testEdrAgentCollection(self):
client_id = db_test_utils.InitializeClient(data_store.REL_DB)
artifact_source = rdf_artifacts.ArtifactSource()
artifact_source.type = rdf_artifacts.ArtifactSource.SourceType.COMMAND
artifact_source.attributes = {"cmd": "/bin/echo", "args": ["1337"]}
artifact = rdf_artifacts.Artifact()
artifact.name = "Foo"
artifact.doc = "Lorem ipsum."
artifact.sources = [artifact_source]
class FooParser(parsers.SingleResponseParser):
supported_artifacts = ["Foo"]
def ParseResponse(
self,
knowledge_base: rdf_client.KnowledgeBase,
response: rdf_client_action.ExecuteResponse,
) -> Iterator[rdf_client.EdrAgent]:
edr_agent = rdf_client.EdrAgent()
edr_agent.name = "echo"
edr_agent.agent_id = response.stdout.decode("utf-8")
yield edr_agent
class EchoActionMock(action_mocks.InterrogatedClient):
def ExecuteCommand(
self,
args: rdf_client_action.ExecuteRequest,
) -> Iterable[rdf_client_action.ExecuteResponse]:
response = rdf_client_action.ExecuteResponse()
response.stdout = " ".join(args.args).encode("utf-8")
response.exit_status = 0
return [response]
with mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry()) as registry:
registry.RegisterArtifact(artifact)
with test_lib.ConfigOverrider({"Artifacts.edr_agents": ["Foo"]}):
with parser_test_lib._ParserContext("Foo", FooParser):
flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock=EchoActionMock(),
client_id=client_id,
creator=self.test_username)
flow_test_lib.FinishAllFlowsOnClient(client_id)
snapshot = data_store.REL_DB.ReadClientSnapshot(client_id)
self.assertLen(snapshot.edr_agents, 1)
self.assertEqual(snapshot.edr_agents[0].name, "echo")
self.assertEqual(snapshot.edr_agents[0].agent_id, "1337")
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 testMultiFileSuccess(self):
class ThudParser(parsers.MultiFileParser):
supported_artifacts = ["Thud"]
def ParseFiles(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspecs: Collection[rdf_paths.PathSpec],
filedescs: Collection[file_store.BlobStream],
) -> Iterable[rdf_protodict.Dict]:
results = []
for pathspec, filedesc in zip(pathspecs, filedescs):
result = rdf_protodict.Dict()
result["path"] = pathspec.path
result["content"] = filedesc.Read()
results.append(result)
return results
with parser_test_lib._ParserContext("Thud", ThudParser):
factory = parsers.ArtifactParserFactory("Thud")
client_id = self.client_id
knowledge_base = rdf_client.KnowledgeBase()
stat_entry_foo = rdf_client_fs.StatEntry()
stat_entry_foo.pathspec.path = "quux/foo"
stat_entry_foo.pathspec.pathtype = rdf_paths.PathSpec.PathType.OS
self._WriteFile(stat_entry_foo.pathspec.path, b"FOO")
stat_entry_bar = rdf_client_fs.StatEntry()
stat_entry_bar.pathspec.path = "quux/bar"
stat_entry_bar.pathspec.pathtype = rdf_paths.PathSpec.PathType.OS
self._WriteFile(stat_entry_bar.pathspec.path, b"BAR")
applicator = artifact.ParserApplicator(factory, client_id,
knowledge_base)
applicator.Apply([stat_entry_foo, stat_entry_bar])
errors = list(applicator.Errors())
self.assertEmpty(errors)
responses = list(applicator.Responses())
self.assertLen(responses, 2)
self.assertEqual(responses[0], {
"path": "quux/foo",
"content": b"FOO"
})
self.assertEqual(responses[1], {
"path": "quux/bar",
"content": b"BAR"
})
def testMultiResponseParserNames(self):
class FooParser(parsers.MultiResponseParser[None]):
supported_artifacts = ["Quux"]
def ParseResponses(
self,
knowledge_base: rdf_client.KnowledgeBase,
responses: Iterable[rdfvalue.RDFValue],
) -> Iterator[None]:
raise NotImplementedError()
with parser_test_lib._ParserContext("Foo", FooParser):
quux_factory = parsers.ArtifactParserFactory("Quux")
self.assertEqual(list(quux_factory.MultiResponseParserNames()), ["Foo"])
def testMultiFileParserNames(self):
class FooParser(parsers.MultiFileParser[None]):
supported_artifacts = ["Quux"]
def ParseFiles(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspecs: Iterable[rdf_paths.PathSpec],
filedescs: Iterable[IO[bytes]],
) -> Iterator[None]:
raise NotImplementedError()
with parser_test_lib._ParserContext("Foo", FooParser):
quux_factory = parsers.ArtifactParserFactory("Quux")
self.assertEqual(list(quux_factory.MultiFileParserNames()), ["Foo"])
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 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 testMultiFileError(self):
class ThudParseError(parsers.ParseError):
pass
class ThudParser(parsers.MultiFileParser):
supported_artifacts = ["Thud"]
def ParseFiles(
self,
knowledge_base: rdf_client.KnowledgeBase,
pathspecs: Collection[rdf_paths.PathSpec],
filedescs: Collection[file_store.BlobStream],
) -> Iterable[rdf_protodict.Dict]:
del knowledge_base, pathspecs, filedescs # Unused.
raise ThudParseError("Lorem ipsum.")
with parser_test_lib._ParserContext("Thud", ThudParser):
factory = parsers.ArtifactParserFactory("Thud")
client_id = self.client_id
knowledge_base = rdf_client.KnowledgeBase()
stat_entry = rdf_client_fs.StatEntry()
stat_entry.pathspec.path = "foo/bar/baz"
stat_entry.pathspec.pathtype = rdf_paths.PathSpec.PathType.OS
self._WriteFile(stat_entry.pathspec.path, b"\xff\x00\xff")
applicator = artifact.ParserApplicator(factory, client_id,
knowledge_base)
applicator.Apply([stat_entry])
errors = list(applicator.Errors())
self.assertLen(errors, 1)
self.assertIsInstance(errors[0], ThudParseError)
responses = list(applicator.Responses())
self.assertEmpty(responses)
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")
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.")
