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