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 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 testCollect(self):
data_store.REL_DB.WriteClientMetadata(
client_id=ClientTest.FAKE_CLIENT_ID, fleetspeak_enabled=False)
client = rdf_objects.ClientSnapshot(
client_id=ClientTest.FAKE_CLIENT_ID)
client.knowledge_base.os = 'test-os'
data_store.REL_DB.WriteClientSnapshot(client)
with mock.patch.object(artifact_registry, 'REGISTRY',
artifact_registry.ArtifactRegistry()):
source = rdf_artifacts.ArtifactSource(
type=artifact_pb2.ArtifactSource.COMMAND,
attributes={
'cmd': '/bin/echo',
'args': ['1']
})
artifact = rdf_artifacts.Artifact(name='FakeArtifact',
sources=[source],
doc='fake artifact doc')
artifact_registry.REGISTRY.RegisterArtifact(artifact)
client = grr_colab.Client.with_id(ClientTest.FAKE_CLIENT_ID)
results = client.collect('FakeArtifact')
self.assertNotEmpty(results)
self.assertEqual(results[0].stdout, b'1\n')
def testArtifactsFromYamlIgnoresDeprecatedFields(self):
registry = ar.ArtifactRegistry()
yaml = textwrap.dedent("""\
name: Foo
doc: Lorem ipsum.
labels: ['bar', 'baz']
sources:
- type: PATH
attributes:
paths: ['/bar', '/baz']
---
name: Quux
doc: Lorem ipsum.
labels: ['norf', 'thud']
sources:
- type: PATH
attributes:
paths: ['/norf', '/thud']
""")
artifacts = registry.ArtifactsFromYaml(yaml)
artifacts.sort(key=lambda artifact: artifact.name)
self.assertLen(artifacts, 2)
self.assertEqual(artifacts[0].name, "Foo")
self.assertFalse(artifacts[0].HasField("labels"))
self.assertEqual(artifacts[1].name, "Quux")
self.assertFalse(artifacts[1].HasField("labels"))
def testNewArtifactLoaded(self):
"""Simulate a new artifact being loaded into the store via the UI."""
cmd_artifact = """name: "TestCmdArtifact"
doc: "Test command artifact for dpkg."
sources:
- type: "COMMAND"
attributes:
cmd: "/usr/bin/dpkg"
args: ["--list"]
labels: [ "Software" ]
supported_os: [ "Linux" ]
"""
no_datastore_artifact = """name: "NotInDatastore"
doc: "Test command artifact for dpkg."
sources:
- type: "COMMAND"
attributes:
cmd: "/usr/bin/dpkg"
args: ["--list"]
labels: [ "Software" ]
supported_os: [ "Linux" ]
"""
test_registry = artifact_registry.ArtifactRegistry()
test_registry.ClearRegistry()
test_registry.AddDatastoreSource(rdfvalue.RDFURN("aff4:/artifact_store"))
test_registry._dirty = False
with utils.Stubber(artifact_registry, "REGISTRY", test_registry):
with self.assertRaises(rdf_artifacts.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("TestCmdArtifact")
with self.assertRaises(rdf_artifacts.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("NotInDatastore")
# Add artifact to datastore but not registry
artifact_coll = artifact_registry.ArtifactCollection(
rdfvalue.RDFURN("aff4:/artifact_store"))
with data_store.DB.GetMutationPool() as pool:
for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml(
cmd_artifact):
artifact_coll.Add(artifact_val, mutation_pool=pool)
# Add artifact to registry but not datastore
for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml(
no_datastore_artifact):
artifact_registry.REGISTRY.RegisterArtifact(
artifact_val, source="datastore", overwrite_if_exists=False)
# We need to reload all artifacts from the data store before trying to get
# the artifact.
artifact_registry.REGISTRY.ReloadDatastoreArtifacts()
self.assertTrue(artifact_registry.REGISTRY.GetArtifact("TestCmdArtifact"))
# We registered this artifact with datastore source but didn't
# write it into aff4. This simulates an artifact that was
# uploaded in the UI then later deleted. We expect it to get
# cleared when the artifacts are reloaded from the datastore.
with self.assertRaises(rdf_artifacts.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("NotInDatastore")
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 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 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 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 setUp(self):
super(ArtifactMigrationTest, self).setUp()
self._db_patcher = mock.patch.object(
data_store, "REL_DB",
db.DatabaseValidationWrapper(mem.InMemoryDB()))
self._db_patcher.start()
self._artifact_patcher = mock.patch.object(
artifact_registry, "REGISTRY",
artifact_registry.ArtifactRegistry())
self._artifact_patcher.start()
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 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 CreateTestArtifactRegistry():
r = artifact_registry.ArtifactRegistry()
r.AddDirSource(os.path.join(config.CONFIG["Test.data_dir"], "artifacts"))
r.AddDatastoreSources([aff4.ROOT_URN.Add("artifact_store")])
return r
def CreateDatastoreOnlyArtifactRegistry():
r = artifact_registry.ArtifactRegistry()
r.AddDatastoreSources([aff4.ROOT_URN.Add("artifact_store")])
return r
def CreateDefaultArtifactRegistry():
r = artifact_registry.ArtifactRegistry()
r.AddDefaultSources()
return r
def CreateTestArtifactRegistry(): r = artifact_registry.ArtifactRegistry() r.AddDirSource(os.path.join(config.CONFIG["Test.data_dir"], "artifacts")) return r
def CreateDatastoreOnlyArtifactRegistry(): return artifact_registry.ArtifactRegistry()
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")
