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)