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