def _ScheduleCollection(self):
# Schedule any new artifacts for which we have now fulfilled dependencies.
for artifact_name in self.state.awaiting_deps_artifacts:
artifact_obj = artifact_registry.REGISTRY.GetArtifact(
artifact_name)
deps = artifact_registry.GetArtifactPathDependencies(artifact_obj)
if set(deps).issubset(self.state.fulfilled_deps):
self.state.in_flight_artifacts.add(artifact_name)
self.state.awaiting_deps_artifacts.remove(artifact_name)
self.CallFlow(
# TODO(user): dependency loop with flows/general/collectors.py.
# collectors.ArtifactCollectorFlow.__name__,
"ArtifactCollectorFlow",
artifact_list=[artifact_name],
next_state=compatibility.GetName(self.ProcessBase),
request_data={"artifact_name": artifact_name},
knowledge_base=self.state.knowledge_base)
# If we're not done but not collecting anything, start accepting the partial
# dependencies as full, and see if we can complete.
if (self.state.awaiting_deps_artifacts
and not self.state.in_flight_artifacts):
if self.state.partial_fulfilled_deps:
partial = self.state.partial_fulfilled_deps.pop()
self.Log("Accepting partially fulfilled dependency: %s",
partial)
self.state.fulfilled_deps.add(partial)
self._ScheduleCollection()
def testGetArtifactPathDependencies(self):
sources = [{
"type": rdf_artifacts.ArtifactSource.SourceType.REGISTRY_KEY,
"attributes": {
"keys": [
r"%%current_control_set%%\Control\Session "
r"Manager\Environment\Path"
]
}
},
{
"type": rdf_artifacts.ArtifactSource.SourceType.WMI,
"attributes": {
"query": "SELECT * FROM Win32_UserProfile "
"WHERE SID='%%users.sid%%'"
}
},
{
"type": rdf_artifacts.ArtifactSource.SourceType.GREP,
"attributes": {
"content_regex_list": ["^%%users.username%%:"]
}
}]
artifact = rdf_artifacts.Artifact(
name="artifact",
doc="Doco",
provides=["environ_windir"],
supported_os=["Windows"],
urls=["http://blah"],
sources=sources)
self.assertCountEqual(
[x["type"] for x in artifact.ToPrimitiveDict()["sources"]],
["REGISTRY_KEY", "WMI", "GREP"])
class Parser1(object):
knowledgebase_dependencies = ["appdata", "sid"]
class Parser2(object):
knowledgebase_dependencies = ["sid", "desktop"]
@classmethod
def MockGetClassesByArtifact(unused_cls, _):
return [Parser1, Parser2]
with utils.Stubber(parser.Parser, "GetClassesByArtifact",
MockGetClassesByArtifact):
self.assertCountEqual(
ar.GetArtifactPathDependencies(artifact), [
"appdata", "sid", "desktop", "current_control_set", "users.sid",
"users.username"
])
def testGetArtifacts(self, registry):
registry.AddFileSource(self.test_artifacts_file)
results = registry.GetArtifacts(os_name="Windows")
for result in results:
self.assertTrue("Windows" in result.supported_os
or not result.supported_os)
results = registry.GetArtifacts(
os_name="Windows",
name_list=["TestAggregationArtifact", "TestFileArtifact"])
# TestFileArtifact doesn't match the OS criteria
self.assertCountEqual([x.name for x in results],
["TestAggregationArtifact"])
for result in results:
self.assertTrue("Windows" in result.supported_os
or not result.supported_os)
results = registry.GetArtifacts(
os_name="Windows",
source_type=rdf_artifacts.ArtifactSource.SourceType.REGISTRY_VALUE,
name_list=["DepsProvidesMultiple"])
self.assertEqual(results.pop().name, "DepsProvidesMultiple")
# Check supported_os = [] matches any OS
results = registry.GetArtifacts(os_name="Windows",
name_list=["TestRegistryKey"])
self.assertLen(results, 1)
artifact = results.pop()
self.assertEqual(artifact.name, "TestRegistryKey")
self.assertEqual(artifact.supported_os, [])
results = registry.GetArtifacts(os_name="Windows",
exclude_dependents=True)
for result in results:
self.assertFalse(ar.GetArtifactPathDependencies(result))
# Check provides filtering
results = registry.GetArtifacts(os_name="Windows",
provides=["users.homedir", "domain"])
for result in results:
# provides contains at least one of the filter strings
self.assertGreaterEqual(
len(
set(result.provides).union(set(["users.homedir",
"domain"]))), 1)
results = registry.GetArtifacts(os_name="Windows",
provides=["nothingprovidesthis"])
self.assertEmpty(results)
def testPrepareArtifactsWithKBInitialization(self):
"""Test the preparation of artifacts for the KB initialization."""
artifact_list = ["TestFilesArtifact", "DepsWindir"]
self.SetOS("Windows")
recollect_knowledge_base = True
args = self.ArtifactCollectorArgs(artifact_list, recollect_knowledge_base)
self.assertLen(args.artifacts, 2)
artifact_names = [str(a.name) for a in args.artifacts]
self.assertEqual(artifact_names, ["DepsControlSet", "DepsWindir"])
first_artifact = artifact_registry.REGISTRY.GetArtifact(artifact_list[0])
dependencies = artifact_registry.GetArtifactPathDependencies(first_artifact)
self.assertEqual(dependencies, set([]))
def _AddDependencyEdges(self, rdf_artifact):
"""Add an edge for every dependency of the given artifact.
This method gets the attribute names for a given artifact and for every
attribute it adds a directed edge from the attribute node to the artifact
node. If an artifact does not have any dependencies it is added to the set
of reachable nodes.
Args:
rdf_artifact: The artifact object.
"""
artifact_dependencies = artifact_registry.GetArtifactPathDependencies(
rdf_artifact)
if artifact_dependencies:
for attribute in artifact_dependencies:
self._AddEdge(attribute, rdf_artifact.name)
else:
self.reachable_nodes.add(rdf_artifact.name)
self.graph[rdf_artifact.name].is_provided = True
def testGetArtifactPathDependencies(self):
sources = [{
"type": rdf_artifacts.ArtifactSource.SourceType.REGISTRY_KEY,
"attributes": {
"keys": [
r"%%current_control_set%%\Control\Session "
r"Manager\Environment\Path"
]
}
},
{
"type": rdf_artifacts.ArtifactSource.SourceType.WMI,
"attributes": {
"query": "SELECT * FROM Win32_UserProfile "
"WHERE SID='%%users.sid%%'"
}
},
{
"type": rdf_artifacts.ArtifactSource.SourceType.GREP,
"attributes": {
"content_regex_list": ["^%%users.username%%:"]
}
}]
artifact = rdf_artifacts.Artifact(
name="artifact",
doc="Doco",
provides=["environ_windir"],
supported_os=["Windows"],
urls=["http://blah"],
sources=sources)
self.assertCountEqual(
[x["type"] for x in artifact.ToPrimitiveDict()["sources"]],
["REGISTRY_KEY", "WMI", "GREP"])
self.assertCountEqual(
ar.GetArtifactPathDependencies(artifact), [
"appdata", "sid", "desktop", "current_control_set", "users.sid",
"users.username"
])
def BuildArtifactDescriptors(self, artifacts_list):
result = []
for artifact_val in artifacts_list:
descriptor = rdf_artifacts.ArtifactDescriptor(
artifact=artifact_val,
dependencies=sorted(
artifact_registry.GetArtifactDependencies(artifact_val)),
path_dependencies=sorted(
artifact_registry.GetArtifactPathDependencies(artifact_val)),
error_message=artifact_val.error_message,
is_custom=artifact_val.loaded_from.startswith("datastore:"))
factory = parsers.ArtifactParserFactory(str(artifact_val.name))
for parser in factory.AllParsers():
parser_cls = type(parser)
descriptor.processors.append(
rdf_artifacts.ArtifactProcessorDescriptor.FromParser(parser_cls))
result.append(descriptor)
return result
def BuildArtifactDescriptors(self, artifacts_list):
result = []
for artifact_val in artifacts_list:
descriptor = rdf_artifacts.ArtifactDescriptor(
artifact=artifact_val,
dependencies=sorted(
artifact_registry.GetArtifactDependencies(artifact_val)),
path_dependencies=sorted(
artifact_registry.GetArtifactPathDependencies(artifact_val)),
error_message=artifact_val.error_message,
is_custom=artifact_val.loaded_from.startswith("datastore:"))
for processor in parser.Parser.GetClassesByArtifact(artifact_val.name):
descriptor.processors.append(
rdf_artifacts.ArtifactProcessorDescriptor(
name=processor.__name__,
output_types=processor.output_types,
description=processor.GetDescription()))
result.append(descriptor)
return result