def testUploadArtifactYamlFileAndDumpToYaml(self):
artifact_registry.REGISTRY.ClearRegistry()
artifact_registry.REGISTRY.ClearSources()
artifact_registry.REGISTRY._CheckDirty()
try:
test_artifacts_file = os.path.join(config.CONFIG["Test.data_dir"],
"artifacts",
"test_artifacts.json")
with io.open(test_artifacts_file, mode="r",
encoding="utf-8") as filedesc:
artifact.UploadArtifactYamlFile(filedesc.read())
loaded_artifacts = artifact_registry.REGISTRY.GetArtifacts()
self.assertGreaterEqual(len(loaded_artifacts), 20)
self.assertIn("DepsWindirRegex",
[a.name for a in loaded_artifacts])
# Now dump back to YAML.
yaml_data = artifact_registry.REGISTRY.DumpArtifactsToYaml()
for snippet in [
"name: TestFilesArtifact",
"urls:\\s*- https://msdn.microsoft.com/en-us/library/",
"returned_types:\\s*- SoftwarePackage",
"args:\\s*- --list",
"cmd: /usr/bin/dpkg",
]:
self.assertRegexpMatches(yaml_data, snippet)
finally:
artifact.ArtifactLoader().RunOnce()
def testUploadArtifactYamlFileAndDumpToYaml(self):
artifact_registry.REGISTRY.ClearRegistry()
artifact_registry.REGISTRY.ClearSources()
artifact_registry.REGISTRY._CheckDirty()
try:
test_artifacts_file = os.path.join(config.CONFIG["Test.data_dir"],
"artifacts",
"test_artifacts.json")
filecontent = open(test_artifacts_file, "rb").read()
artifact.UploadArtifactYamlFile(filecontent)
loaded_artifacts = artifact_registry.REGISTRY.GetArtifacts()
self.assertGreaterEqual(len(loaded_artifacts), 20)
self.assertIn("DepsWindirRegex",
[a.name for a in loaded_artifacts])
# Now dump back to YAML.
yaml_data = artifact_registry.REGISTRY.DumpArtifactsToYaml()
for snippet in [
"name: TestFilesArtifact",
"urls: ['https://msdn.microsoft.com/en-us/library/aa384749%28v=vs.85",
"returned_types: [SoftwarePackage]",
"args: [--list]",
"cmd: /usr/bin/dpkg",
]:
self.assertIn(snippet, yaml_data)
finally:
artifact.ArtifactLoader().RunOnce()
def testGetKBDependencies(self):
"""Test that KB dependencies are calculated correctly."""
artifact_registry.REGISTRY.ClearSources()
try:
test_artifacts_file = os.path.join(config.CONFIG["Test.data_dir"],
"artifacts",
"test_artifacts.json")
artifact_registry.REGISTRY.AddFileSource(test_artifacts_file)
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base":
[
"DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindir",
"DepsWindirRegex", "DepsControlSet", "FakeArtifact"
],
"Artifacts.knowledge_base_additions": ["DepsHomedir2"],
"Artifacts.knowledge_base_skip": ["DepsWindir"],
"Artifacts.knowledge_base_heavyweight": ["FakeArtifact"]
}):
args = artifact.KnowledgeBaseInitializationArgs(
lightweight=True)
kb_init = aff4_flows.KnowledgeBaseInitializationFlow(
None, token=self.token)
kb_init.args = args
kb_init.state["all_deps"] = set()
kb_init.state["awaiting_deps_artifacts"] = []
kb_init.state["knowledge_base"] = rdf_client.KnowledgeBase(
os="Windows")
no_deps = kb_init.GetFirstFlowsForCollection()
self.assertCountEqual(no_deps,
["DepsControlSet", "DepsHomedir2"])
self.assertCountEqual(kb_init.state.all_deps, [
"users.homedir", "users.desktop", "users.username",
"environ_windir", "current_control_set"
])
self.assertCountEqual(kb_init.state.awaiting_deps_artifacts, [
"DepsParent", "DepsDesktop", "DepsHomedir",
"DepsWindirRegex"
])
finally:
artifact.ArtifactLoader().RunOnce()
