def testPrepareArtifactFilesBundle(self): """Test the preparation of ArtifactFiles Args.""" artifact_list = ["TestArtifactFilesArtifact"] kb = rdf_client.KnowledgeBase() kb.os = "Linux" file_path = os.path.join(self.base_path, "numbers.txt") source = rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.FILE, attributes={"paths": [file_path]}) artifact_obj = artifact_registry.REGISTRY.GetArtifact("TestFileArtifact") artifact_obj.sources.append(source) artifact_bundle = collectors.GetArtifactCollectorArgs(kb, artifact_list) art_obj = artifact_bundle.artifacts[0] self.assertEqual(art_obj.name, "TestArtifactFilesArtifact") source = art_obj.sources[0] self.assertEqual(source.base_source.type, "ARTIFACT_FILES") sub_artifact_source = source.artifact_sources[0] self.assertEqual(sub_artifact_source.base_source.type, "FILE")
def testKnowledgeBase(self): kb = rdf_client.KnowledgeBase() kb.os = "Windows" artifact_bundle = collectors.GetArtifactCollectorArgs(kb, set(), []) self.assertEqual(artifact_bundle.knowledge_base.os, "Windows")
def ArtifactCollectorArgs(self, artifact_list, collect_knowledge_base=False): flow_args = rdf_artifacts.ArtifactCollectorFlowArgs( artifact_list=artifact_list, recollect_knowledge_base=collect_knowledge_base) return collectors.GetArtifactCollectorArgs(flow_args, self.knowledge_base)
def testDuplicationChecks(self): """Test duplicated artifacts are only processed once.""" artifact_list = [ "TestAggregationArtifact", "TestFilesArtifact", "TestCmdArtifact", "TestFilesArtifact" ] kb = rdf_client.KnowledgeBase() kb.os = "Linux" artifact_bundle = collectors.GetArtifactCollectorArgs(kb, artifact_list) artifacts_objects = list(artifact_bundle.artifacts) self.assertEqual(len(artifacts_objects), 2)
def testPrepareBasicArtifactBundle(self): """Test we can prepare a basic artifact.""" artifact_list = ["TestCmdArtifact"] kb = rdf_client.KnowledgeBase() kb.os = "Linux" artifact_bundle = collectors.GetArtifactCollectorArgs(kb, artifact_list) artifacts_objects = list(artifact_bundle.artifacts) art_obj = artifacts_objects[0] source = list(art_obj.sources)[0] self.assertEqual(art_obj.name, "TestCmdArtifact") self.assertEqual(source.base_source.attributes["cmd"], "/usr/bin/dpkg") self.assertEqual(source.base_source.attributes.get("args", []), ["--list"])
def testPrepareAggregatedArtifactBundle(self): """Test we can prepare the source artifacts of an aggregation artifact.""" artifact_list = ["TestAggregationArtifact"] kb = rdf_client.KnowledgeBase() kb.os = "Windows" artifact_bundle = collectors.GetArtifactCollectorArgs(kb, artifact_list) artifacts_objects = list(artifact_bundle.artifacts) self.assertEqual(len(artifact_bundle.artifacts), 1) art_obj = artifacts_objects[0] self.assertEqual(art_obj.name, "TestAggregationArtifact") self.assertEqual(len(art_obj.sources), 2) source = list(art_obj.sources)[0] self.assertEqual(source.base_source.type, "GRR_CLIENT_ACTION") source = list(art_obj.sources)[1] self.assertEqual(source.base_source.type, "COMMAND")
def testPrepareMultipleArtifacts(self): """Test we can prepare multiple artifacts of different types.""" artifact_list = [ "TestFilesArtifact", "DepsWindirRegex", "DepsProvidesMultiple", "WMIActiveScriptEventConsumer" ] kb = rdf_client.KnowledgeBase() kb.os = "Windows" artifact_bundle = collectors.GetArtifactCollectorArgs(kb, artifact_list) artifacts_objects = list(artifact_bundle.artifacts) self.assertEqual(len(artifacts_objects), 3) self.assertEqual(artifacts_objects[0].name, "DepsWindirRegex") self.assertEqual(artifacts_objects[1].name, "DepsProvidesMultiple") self.assertEqual(artifacts_objects[2].name, "WMIActiveScriptEventConsumer") art_obj = artifacts_objects[2] source = list(art_obj.sources)[0] self.assertEqual(source.base_source.attributes["query"], "SELECT * FROM ActiveScriptEventConsumer")
def ArtifactCollectorArgs(self, artifact_list, collect_knowledge_base=False): return collectors.GetArtifactCollectorArgs( self.knowledge_base, artifact_list, recollect_knowledge_base=collect_knowledge_base)