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)
