def testInterpolateArgs(self): collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token) collect_flow.state.Register("knowledge_base", rdfvalue.KnowledgeBase()) collect_flow.current_artifact_name = "blah" collect_flow.state.knowledge_base.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******")) collect_flow.state.knowledge_base.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******")) test_rdf = rdfvalue.KnowledgeBase() action_args = {"usernames": ["%%users.username%%", "%%users.username%%"], "nointerp": "asdfsdf", "notastring": test_rdf} kwargs = collect_flow.InterpolateDict(action_args) self.assertItemsEqual(kwargs["usernames"], ["test1", "test2", "test1", "test2"]) self.assertEqual(kwargs["nointerp"], "asdfsdf") self.assertEqual(kwargs["notastring"], test_rdf) # We should be using an array since users.username will expand to multiple # values. self.assertRaises(ValueError, collect_flow.InterpolateDict, {"bad": "%%users.username%%"}) list_args = collect_flow.InterpolateList(["%%users.username%%", "%%users.username%%aa"]) self.assertItemsEqual(list_args, ["test1", "test2", "test1aa", "test2aa"]) list_args = collect_flow.InterpolateList(["one"]) self.assertEqual(list_args, ["one"])
def setUp(self): # Set the KB to an empty object client = aff4.FACTORY.Open(self.client_id, mode="rw", token=self.token) self.old_kb = client.Get(client.Schema.KNOWLEDGE_BASE) client.Set(client.Schema.KNOWLEDGE_BASE, rdfvalue.KnowledgeBase()) client.Flush() super(TestParserDependency, self).setUp()
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 testFindsKeyWithInterpolatedGlobWithoutConditions(self): # Initialize client's knowledge base in order for the interpolation # to work. user = rdfvalue.KnowledgeBaseUser( sid="S-1-5-21-2911950750-476812067-1487428992-1001") kb = rdfvalue.KnowledgeBase(users=[user]) with aff4.FACTORY.Open(self.client_id, mode="rw", token=self.token) as client: client.Set(client.Schema.KNOWLEDGE_BASE, kb) self.RunFlow([ "HKEY_USERS/%%users.sid%%/Software/Microsoft/Windows/" "CurrentVersion/*" ]) results = self.GetResults() self.assertEqual(len(results), 1) key = ("/HKEY_USERS/S-1-5-21-2911950750-476812067-1487428992-1001/" "Software/Microsoft/Windows/CurrentVersion/Explorer") self.assertEqual(results[0].stat_entry.aff4path, "aff4:/C.1000000000000000/registry" + key) self.assertEqual(results[0].stat_entry.pathspec.path, key) self.assertEqual(results[0].stat_entry.pathspec.pathtype, rdfvalue.PathSpec.PathType.REGISTRY)
def testInterpolation(self): """Check we can interpolate values from the knowledge base.""" kb = rdfvalue.KnowledgeBase() self.assertRaises( artifact_lib.KnowledgeBaseInterpolationError, list, artifact_lib.InterpolateKbAttributes("test%%users.username%%test", kb)) kb.users.Append(rdfvalue.KnowledgeBaseUser(username="******", uid=1)) kb.users.Append(rdfvalue.KnowledgeBaseUser(username="******", uid=2)) kb.Set("environ_allusersprofile", "c:\\programdata") paths = artifact_lib.InterpolateKbAttributes( "test%%users.username%%test", kb) paths = list(paths) self.assertEqual(len(paths), 2) self.assertItemsEqual(paths, ["testjoetest", "testjimtest"]) paths = artifact_lib.InterpolateKbAttributes( "%%environ_allusersprofile%%\\a", kb) self.assertEqual(list(paths), ["c:\\programdata\\a"]) self.assertRaises( artifact_lib.KnowledgeBaseInterpolationError, list, artifact_lib.InterpolateKbAttributes("%%nonexistent%%\\a", kb)) kb.Set("environ_allusersprofile", "") self.assertRaises( artifact_lib.KnowledgeBaseInterpolationError, list, artifact_lib.InterpolateKbAttributes( "%%environ_allusersprofile%%\\a", kb))
def Validate(self): """Attempt to validate the artifact has been well defined. This is used to enforce Artifact rules. Since it checks all dependencies are present, this method can only be called once all artifacts have been loaded into the registry. Use ValidateSyntax to check syntax for each artifact on import. Raises: ArtifactDefinitionError: If artifact is invalid. """ cls_name = self.name self.ValidateSyntax() # Check all path dependencies exist in the knowledge base. valid_fields = rdfvalue.KnowledgeBase().GetKbFieldNames() for dependency in self.GetArtifactPathDependencies(): if dependency not in valid_fields: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has an invalid path dependency %s. Artifacts must use " "defined knowledge attributes." % (cls_name, dependency)) # Check all artifact dependencies exist. for dependency in self.GetArtifactDependencies(): if dependency not in artifact_lib.ArtifactRegistry.artifacts: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has an invalid dependency %s . Could not find artifact" " definition." % (cls_name, dependency))
def testGrep(self): class MockCallFlow(object): def CallFlow(self, *args, **kwargs): self.args = args self.kwargs = kwargs mock_call_flow = MockCallFlow() with test_lib.Stubber(collectors.ArtifactCollectorFlow, "CallFlow", mock_call_flow.CallFlow): collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token) collect_flow.state.Register("knowledge_base", rdfvalue.KnowledgeBase()) collect_flow.current_artifact_name = "blah" collect_flow.state.knowledge_base.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******")) collect_flow.state.knowledge_base.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******")) collector = rdfvalue.Collector(action="Grep", args={ "path_list": ["/etc/passwd"], "content_regex_list": [r"^a%%users.username%%b$"] }) collect_flow.Grep(collector, rdfvalue.PathSpec.PathType.TSK) filters = mock_call_flow.kwargs["filters"] regexes = [ f.contents_regex_match.regex.SerializeToString() for f in filters ] self.assertItemsEqual(regexes, [r"^atest1b$", r"^atest2b$"]) self.assertEqual(mock_call_flow.kwargs["paths"], ["/etc/passwd"])
def testGrep(self): class MockCallFlow(object): def CallFlow(self, *args, **kwargs): self.args = args self.kwargs = kwargs mock_call_flow = MockCallFlow() with utils.Stubber(collectors.ArtifactCollectorFlow, "CallFlow", mock_call_flow.CallFlow): collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token) collect_flow.state.Register("knowledge_base", rdfvalue.KnowledgeBase()) collect_flow.current_artifact_name = "blah" collect_flow.state.knowledge_base.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******")) collect_flow.state.knowledge_base.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******")) collector = rdfvalue.Collector( collector_type=rdfvalue.Collector.CollectorType.GREP, args={"path_list": ["/etc/passwd"], "content_regex_list": [r"^a%%users.username%%b$"]}) collect_flow.Grep(collector, rdfvalue.PathSpec.PathType.TSK) conditions = mock_call_flow.kwargs["conditions"] self.assertEqual(len(conditions), 1) regexes = conditions[0].contents_regex_match.regex.SerializeToString() self.assertItemsEqual(regexes.split("|"), ["(^atest1b$)", "(^atest2b$)"]) self.assertEqual(mock_call_flow.kwargs["paths"], ["/etc/passwd"])
def testUserMergeWindows(self): """Check Windows users are accurately merged.""" kb = rdfvalue.KnowledgeBase() self.assertEqual(len(kb.users), 0) kb.MergeOrAddUser(rdfvalue.KnowledgeBaseUser(sid="1234")) self.assertEqual(len(kb.users), 1) kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(sid="5678", username="******")) self.assertEqual(len(kb.users), 2) _, conflicts = kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(sid="5678", username="******")) self.assertEqual(len(kb.users), 2) self.assertEqual(conflicts[0], ("username", "test1", "test2")) self.assertEqual(kb.GetUser(sid="5678").username, "test2") # This should merge on user name as we have no other data. kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******", homedir="a")) self.assertEqual(len(kb.users), 2) # This should create a new user since the sid is different. new_attrs, conflicts = kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******", sid="12345", temp="/blah")) self.assertEqual(len(kb.users), 3) self.assertItemsEqual(new_attrs, ["users.username", "users.temp", "users.sid"]) self.assertEqual(conflicts, [])
def testUserMergeLinux(self): """Check Linux users are accurately merged.""" kb = rdfvalue.KnowledgeBase() self.assertEqual(len(kb.users), 0) kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******", last_logon=1111)) self.assertEqual(len(kb.users), 1) # This should merge since the username is the same. kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(uid="12", username="******")) self.assertEqual(len(kb.users), 1) # This should create a new record because the uid is different kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******", uid="13", desktop="/home/blake/Desktop")) self.assertEqual(len(kb.users), 2) kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******", uid="14", desktop="/home/blake/Desktop")) self.assertEqual(len(kb.users), 3) # Check merging where we don't specify uid works new_attrs, conflicts = kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******", desktop="/home/blakey/Desktop")) self.assertEqual(len(kb.users), 3) self.assertItemsEqual(new_attrs, ["users.username", "users.desktop"]) self.assertItemsEqual( conflicts, [("desktop", u"/home/blake/Desktop", u"/home/blakey/Desktop")])
def Layout(self, request, response): self.completions = rdfvalue.KnowledgeBase().GetKbFieldNames() response = super(GlobExpressionFormRenderer, self).Layout(request, response) return self.CallJavascript(response, "GlobExpressionFormRenderer.Layout", prefix=self.prefix, completions=self.completions)
def ValidateSyntax(self): """Validate artifact syntax. This method can be used to validate individual artifacts as they are loaded, without needing all artifacts to be loaded first, as for Validate(). Raises: ArtifactDefinitionError: If artifact is invalid. """ cls_name = self.name if not self.doc: raise artifact_registry.ArtifactDefinitionError( "Artifact %s has missing doc" % cls_name) for supp_os in self.supported_os: if supp_os not in SUPPORTED_OS_LIST: raise artifact_registry.ArtifactDefinitionError( "Artifact %s has invalid supported_os %s" % (cls_name, supp_os)) for condition in self.conditions: try: of = objectfilter.Parser(condition).Parse() of.Compile(objectfilter.BaseFilterImplementation) except ConditionError as e: raise artifact_registry.ArtifactDefinitionError( "Artifact %s has invalid condition %s. %s" % ( cls_name, condition, e)) for label in self.labels: if label not in ARTIFACT_LABELS: raise artifact_registry.ArtifactDefinitionError( "Artifact %s has an invalid label %s. Please use one from " "ARTIFACT_LABELS." % (cls_name, label)) # Anything listed in provides must be defined in the KnowledgeBase valid_provides = rdfvalue.KnowledgeBase().GetKbFieldNames() for kb_var in self.provides: if kb_var not in valid_provides: raise artifact_registry.ArtifactDefinitionError( "Artifact %s has broken provides: '%s' not in KB fields: %s" % ( cls_name, kb_var, valid_provides)) # Any %%blah%% path dependencies must be defined in the KnowledgeBase for dep in self.GetArtifactPathDependencies(): if dep not in valid_provides: raise artifact_registry.ArtifactDefinitionError( "Artifact %s has an invalid path dependency: '%s', not in KB " "fields: %s" % (cls_name, dep, valid_provides)) for source in self.sources: try: source.Validate() except Error as e: raise artifact_registry.ArtifactDefinitionError( "Artifact %s has bad source. %s" % (cls_name, e))
def testKnowledgeBaseRetrievalFailures(self): """Test kb retrieval failure modes.""" client = aff4.FACTORY.Open(self.client_id, token=self.token, mode="rw") self.assertRaises(artifact_lib.KnowledgeBaseUninitializedError, artifact.GetArtifactKnowledgeBase, client) kb = rdfvalue.KnowledgeBase() kb.hostname = "test" client.Set(client.Schema.KNOWLEDGE_BASE(kb)) client.Flush(sync=True) self.assertRaises(artifact_lib.KnowledgeBaseAttributesMissingError, artifact.GetArtifactKnowledgeBase, client)
def InitializeKnowledgeBase(self): """Get the existing KB or create a new one if none exists.""" self.client = aff4.FACTORY.Open(self.client_id, token=self.token) # Always create a new KB to override any old values. self.state.knowledge_base = rdfvalue.KnowledgeBase() SetCoreGRRKnowledgeBaseValues(self.state.knowledge_base, self.client) if not self.state.knowledge_base.os: # If we don't know what OS this is, there is no way to proceed. raise flow.FlowError("Client OS not set for: %s, cannot initialize" " KnowledgeBase" % self.client_id)
def SetKnowledgeBase(self, hostname="test.example.com", host_os="Linux", host_data=None): if not host_data: host_data = {} kb = rdfvalue.KnowledgeBase() kb.hostname = hostname kb.os = host_os host_data["KnowledgeBase"] = kb return host_data
def testInterpolation(self): """Check we can interpolate values from the knowledge base.""" kb = rdfvalue.KnowledgeBase() # No users yet, this should raise self.assertRaises( artifact_lib.KnowledgeBaseInterpolationError, list, artifact_lib.InterpolateKbAttributes("test%%users.username%%test", kb)) # Now we have two users kb.users.Append(rdfvalue.KnowledgeBaseUser(username="******", uid=1)) kb.users.Append(rdfvalue.KnowledgeBaseUser(username="******", uid=2)) kb.Set("environ_allusersprofile", "c:\\programdata") paths = artifact_lib.InterpolateKbAttributes( "test%%users.username%%test", kb) paths = list(paths) self.assertEqual(len(paths), 2) self.assertItemsEqual(paths, ["testjoetest", "testjimtest"]) paths = artifact_lib.InterpolateKbAttributes( "%%environ_allusersprofile%%\\a", kb) self.assertEqual(list(paths), ["c:\\programdata\\a"]) # Check a bad attribute raises self.assertRaises( artifact_lib.KnowledgeBaseInterpolationError, list, artifact_lib.InterpolateKbAttributes("%%nonexistent%%\\a", kb)) # Empty values should also raise kb.Set("environ_allusersprofile", "") self.assertRaises( artifact_lib.KnowledgeBaseInterpolationError, list, artifact_lib.InterpolateKbAttributes( "%%environ_allusersprofile%%\\a", kb)) # No users have temp defined, so this should raise self.assertRaises( artifact_lib.KnowledgeBaseInterpolationError, list, artifact_lib.InterpolateKbAttributes("%%users.temp%%\\a", kb)) # One user has users.temp defined, the others do not. This is common on # windows where users have been created but have never logged in. We should # get just one value back. kb.users.Append( rdfvalue.KnowledgeBaseUser( username="******", uid=1, temp="C:\\Users\\jason\\AppData\\Local\\Temp")) paths = artifact_lib.InterpolateKbAttributes(r"%%users.temp%%\abcd", kb) self.assertItemsEqual(paths, ["C:\\Users\\jason\\AppData\\Local\\Temp\\abcd"])
def testParse(self): parser = windows_persistence.WindowsPersistenceMechanismsParser() path = (r"HKEY_LOCAL_MACHINE\Software\Microsoft\Windows\CurrentVersion" r"\Run\test") pathspec = rdfvalue.PathSpec( path=path, pathtype=rdfvalue.PathSpec.PathType.REGISTRY) reg_data = "C:\\blah\\some.exe /v" reg_type = rdfvalue.StatEntry.RegistryType.REG_SZ stat = rdfvalue.StatEntry( aff4path="aff4:/asdfasdf/", pathspec=pathspec, registry_type=reg_type, registry_data=rdfvalue.DataBlob(string=reg_data)) persistence = [stat] image_paths = [ "system32\\drivers\\ACPI.sys", "%systemroot%\\system32\\svchost.exe -k netsvcs", "\\SystemRoot\\system32\\drivers\\acpipmi.sys" ] reg_key = rdfvalue.RDFURN("aff4:/C.1000000000000000/registry" "/HKEY_LOCAL_MACHINE/SYSTEM/ControlSet001" "/services/AcpiPmi") for path in image_paths: serv_info = rdfvalue.WindowsServiceInformation( name="blah", display_name="GRRservice", image_path=path, registry_key=reg_key) persistence.append(serv_info) knowledge_base = rdfvalue.KnowledgeBase() knowledge_base.environ_systemroot = "C:\\Windows" expected = [ "C:\\blah\\some.exe", "C:\\Windows\\system32\\drivers\\ACPI.sys", "C:\\Windows\\system32\\svchost.exe", "C:\\Windows\\system32\\drivers\\acpipmi.sys" ] for index, item in enumerate(persistence): results = list( parser.Parse(item, knowledge_base, rdfvalue.PathSpec.PathType.OS)) self.assertEqual(results[0].pathspec.path, expected[index]) self.assertEqual(len(results), 1)
def testBasicParsing(self): knowledge_base = rdfvalue.KnowledgeBase() knowledge_base.environ_systemdrive = "C:" parser = volatility_artifact_parser.VolatilityVADParser() volatility_data = self.GenerateVADVolatilityResult( ["\\WINDOWS\\system.exe", "\\PROGRAM~1\\PROGRAM\\process.exe"]) expected = [ rdfvalue.PathSpec(path="C:\\WINDOWS\\system.exe", pathtype=rdfvalue.PathSpec.PathType.OS), rdfvalue.PathSpec(path="C:\\PROGRAM~1\\PROGRAM\\process.exe", pathtype=rdfvalue.PathSpec.PathType.OS) ] results = list(parser.Parse(volatility_data, knowledge_base)) self.assertListEqual(results, expected)
def testBasicParsing(self): ps_list_file = os.path.join(config_lib.CONFIG["Test.data_dir"], "rekall_vad_result.dat") result = rdfvalue.RekallResponse( json_messages=open(ps_list_file).read(10000000), plugin="pslist", ) knowledge_base = rdfvalue.KnowledgeBase() knowledge_base.environ_systemdrive = "C:" parser = rekall_artifact_parser.RekallVADParser() parsed_pathspecs = list(parser.Parse(result, knowledge_base)) paths = [p.path for p in parsed_pathspecs] self.assertIn(u"C:\\Windows\\System32\\spoolsv.exe", paths)
def Start(self): """For each artifact, create subflows for each collector.""" self.client = aff4.FACTORY.Open(self.client_id, token=self.token) kb = rdfvalue.KnowledgeBase() SetCoreGRRKnowledgeBaseValues(kb, self.client) if not kb.os: raise flow.FlowError("Client OS not set for: %s, cannot initialize" " KnowledgeBase" % self.client_id) self.state.Register("knowledge_base", kb) self.state.Register("fulfilled_deps", []) self.state.Register("partial_fulfilled_deps", set()) self.state.Register("all_deps", set()) self.state.Register("in_flight_artifacts", []) self.state.Register("awaiting_deps_artifacts", []) self.state.Register("completed_artifacts", []) self.CallFlow("BootStrapKnowledgeBaseFlow", next_state="ProcessBootstrap")
def setUp(self): super(CheckRegistryTests, self).setUp() if self.sw_chk is None: self.sw_chk = _LoadCheck("sw.yaml", "SW-CHECK") checks.CheckRegistry.RegisterCheck(check=self.sw_chk, source="dpkg.out", overwrite_if_exists=True) if self.sshd_chk is None: self.sshd_chk = _LoadCheck("sshd.yaml", "SSHD-CHECK") checks.CheckRegistry.RegisterCheck(check=self.sshd_chk, source="sshd_config", overwrite_if_exists=True) self.kb = rdfvalue.KnowledgeBase() self.kb.hostname = "test.example.com" self.host_data = {"KnowledgeBase": self.kb, "WMIInstalledSoftware": WMI_SW, "DebianPackagesStatus": DPKG_SW, "SshdConfigFile": SSHD_CFG}
def testUserMerge(self): """Check users are accurately merged.""" kb = rdfvalue.KnowledgeBase() self.assertEquals(len(kb.users), 0) kb.MergeOrAddUser(rdfvalue.KnowledgeBaseUser(sid="1234")) self.assertEquals(len(kb.users), 1) kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(sid="5678", username="******")) self.assertEquals(len(kb.users), 2) _, conflicts = kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(sid="5678", username="******")) self.assertEquals(len(kb.users), 2) self.assertEquals(conflicts[0], ("username", "test1", "test2")) self.assertEquals(kb.GetUser(sid="5678").username, "test2") # This should merge on user name as we have no other data. kb.MergeOrAddUser( rdfvalue.KnowledgeBaseUser(username="******", homedir="a")) self.assertEquals(len(kb.users), 2)
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 Validate(self): """Attempt to validate the artifact has been well defined. This is used to enforce Artifact rules. Raises: ArtifactDefinitionError: If artifact is invalid. """ cls_name = self.name if not self.doc: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has missing doc" % cls_name) for supp_os in self.supported_os: if supp_os not in artifact_lib.SUPPORTED_OS_LIST: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has invalid supported_os %s" % (cls_name, supp_os)) for condition in self.conditions: try: of = objectfilter.Parser(condition).Parse() of.Compile(objectfilter.BaseFilterImplementation) except artifact_lib.ConditionError as e: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has invalid condition %s. %s" % (cls_name, condition, e)) for label in self.labels: if label not in artifact_lib.ARTIFACT_LABELS: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has an invalid label %s. Please use one from " "ARTIFACT_LABELS." % (cls_name, label)) # Check all path dependencies exist in the knowledge base. valid_fields = rdfvalue.KnowledgeBase().GetKbFieldNames() for dependency in self.GetArtifactPathDependencies(): if dependency not in valid_fields: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has an invalid path dependency %s. Artifacts must use " "defined knowledge attributes." % (cls_name, dependency)) # Check all artifact dependencies exist. for dependency in self.GetArtifactDependencies(): if dependency not in artifact_lib.ArtifactRegistry.artifacts: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has an invalid dependency %s . Could not find artifact" " definition." % (cls_name, dependency)) for collector in self.collectors: try: collector.Validate() except artifact_lib.Error as e: raise artifact_lib.ArtifactDefinitionError( "Artifact %s has bad collector. %s" % (cls_name, e)) for kb_var in self.provides: if len(kb_var) < 3: # Someone probably interpreted string as list. raise artifact_lib.ArtifactDefinitionError( "Artifact %s has broken provides. %s" % (cls_name, self.provides))
def SetKnowledgeBase(self, hostname, host_os, host_data): kb = rdfvalue.KnowledgeBase() kb.hostname = hostname kb.os = host_os host_data["KnowledgeBase"] = kb