def testSearchDependencies(self): with utils.Stubber(artifact_registry.ArtifactRegistry, "artifacts", {}): # Just use the test artifacts to verify dependency correctness so we # aren't subject to changing dependencies in the whole set test_artifacts_file = os.path.join( config_lib.CONFIG["Test.data_dir"], "test_artifacts.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file]) names, expansions = artifact_registry.ArtifactRegistry.SearchDependencies( "Windows", [u"TestAggregationArtifactDeps", u"DepsParent"]) # This list contains all artifacts that can provide the dependency, e.g. # DepsHomedir and DepsHomedir2 both provide # users.homedir. self.assertItemsEqual(names, [ u"DepsHomedir", u"DepsHomedir2", u"DepsDesktop", u"DepsParent", u"DepsWindir", u"DepsWindirRegex", u"DepsControlSet", u"TestAggregationArtifactDeps" ]) self.assertItemsEqual(expansions, [ "current_control_set", "users.homedir", "users.desktop", "environ_windir", "users.username" ]) # None of these match the OS, so we should get an empty list. names, expansions = artifact_registry.ArtifactRegistry.SearchDependencies( "Darwin", [u"TestCmdArtifact", u"TestFileArtifact"]) self.assertItemsEqual(names, [])
def Run(self): artifact_registry.ArtifactRegistry.ClearRegistry() test_artifacts_file = os.path.join( config_lib.CONFIG["Test.data_dir"], "test_artifact.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file]) self.Check("GET", "/api/artifacts")
def testGetDependencies(self): """Test that dependencies are calculated correctly.""" self.SetupMocks() with test_lib.Stubber(artifact_lib.ArtifactRegistry, "artifacts", {}): test_artifacts_file = os.path.join( config_lib.CONFIG["Test.data_dir"], "test_artifacts.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file]) config_lib.CONFIG.Set("Artifacts.knowledge_base", ["DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindir", "DepsWindirRegex", "DepsControlSet"]) config_lib.CONFIG.Set("Artifacts.knowledge_base_additions", ["DepsHomedir2"]) config_lib.CONFIG.Set("Artifacts.knowledge_base_skip", ["DepsWindir"]) kb_init = artifact.KnowledgeBaseInitializationFlow(None, token=self.token) kb_init.state.Register("all_deps", set()) kb_init.state.Register("awaiting_deps_artifacts", []) kb_init.state.Register("knowledge_base", rdfvalue.KnowledgeBase(os="Windows")) no_deps, all_deps, waiting = kb_init._GetDependencies() self.assertItemsEqual(no_deps, ["DepsControlSet", "DepsHomedir2"]) self.assertItemsEqual(all_deps, ["users.homedir", "users.desktop", "users.username", "environ_windir", "current_control_set"]) self.assertItemsEqual(waiting, ["DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindirRegex"])
def testGetDependencies(self): """Test that dependencies are calculated correctly.""" self.SetupWindowsMocks() with utils.Stubber(artifact_registry.ArtifactRegistry, "artifacts", {}): test_artifacts_file = os.path.join( config_lib.CONFIG["Test.data_dir"], "test_artifacts.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file]) # No dependencies args = artifact.CollectArtifactDependenciesArgs( artifact_list=["DepsHomedir2"]) collect_obj = artifact.CollectArtifactDependencies( None, token=self.token) collect_obj.args = args collect_obj.knowledge_base = None collect_obj.state.Register("all_deps", set()) collect_obj.state.Register("awaiting_deps_artifacts", []) collect_obj.state.Register("knowledge_base", rdfvalue.KnowledgeBase(os="Windows")) no_deps = collect_obj.GetFirstFlowsForCollection() self.assertItemsEqual(no_deps, []) self.assertItemsEqual(collect_obj.state.all_deps, []) self.assertItemsEqual(collect_obj.state.awaiting_deps_artifacts, []) # Dependency tree with a single starting point args = artifact.CollectArtifactDependenciesArgs( artifact_list=["DepsHomedir"]) collect_obj.args = args no_deps = collect_obj.GetFirstFlowsForCollection() self.assertItemsEqual(no_deps, ["DepsControlSet"]) self.assertItemsEqual( collect_obj.state.all_deps, ["environ_windir", "users.username", "current_control_set"]) self.assertItemsEqual(collect_obj.state.awaiting_deps_artifacts, ["DepsWindir", "DepsWindirRegex"])
def testGetKBDependencies(self): """Test that KB dependencies are calculated correctly.""" self.SetupWindowsMocks() with utils.Stubber(artifact_registry.ArtifactRegistry, "artifacts", {}): test_artifacts_file = os.path.join( config_lib.CONFIG["Test.data_dir"], "test_artifacts.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file]) config_lib.CONFIG.Set("Artifacts.knowledge_base", [ "DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindir", "DepsWindirRegex", "DepsControlSet", "FakeArtifact" ]) config_lib.CONFIG.Set("Artifacts.knowledge_base_additions", ["DepsHomedir2"]) config_lib.CONFIG.Set("Artifacts.knowledge_base_skip", ["DepsWindir"]) config_lib.CONFIG.Set("Artifacts.knowledge_base_heavyweight", ["FakeArtifact"]) args = rdfvalue.KnowledgeBaseInitializationArgs(lightweight=True) kb_init = artifact.KnowledgeBaseInitializationFlow( None, token=self.token) kb_init.args = args kb_init.state.Register("all_deps", set()) kb_init.state.Register("awaiting_deps_artifacts", []) kb_init.state.Register("knowledge_base", rdfvalue.KnowledgeBase(os="Windows")) no_deps = kb_init.GetFirstFlowsForCollection() self.assertItemsEqual(no_deps, ["DepsControlSet", "DepsHomedir2"]) self.assertItemsEqual(kb_init.state.all_deps, [ "users.homedir", "users.desktop", "users.username", "environ_windir", "current_control_set" ]) self.assertItemsEqual(kb_init.state.awaiting_deps_artifacts, [ "DepsParent", "DepsDesktop", "DepsHomedir", "DepsWindirRegex" ])
def LoadTestArtifacts(cls): test_artifacts_file = os.path.join(config_lib.CONFIG["Test.data_dir"], "test_artifacts.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file])
def setUp(self): super(ArtifactHandlingTest, self).setUp() test_artifacts_file = os.path.join(config_lib.CONFIG["Test.data_dir"], "test_artifacts.json") artifact_lib.LoadArtifactsFromFiles([test_artifacts_file])