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")