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