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.args = mock.Mock() collect_flow.args.ignore_interpolation_errors = False kb = rdf_client.KnowledgeBase() kb.MergeOrAddUser(rdf_client.User(username="******")) kb.MergeOrAddUser(rdf_client.User(username="******")) collect_flow.state["knowledge_base"] = kb collect_flow.current_artifact_name = "blah" collector = rdf_artifacts.ArtifactSource( type=rdf_artifacts.ArtifactSource.SourceType.GREP, attributes={ "paths": ["/etc/passwd"], "content_regex_list": [r"^a%%users.username%%b$"] }) collect_flow.Grep(collector, rdf_paths.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 testGrepRegexCombination(self): collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token) self.assertEqual(collect_flow._CombineRegex([r"simple"]), "simple") self.assertEqual(collect_flow._CombineRegex(["a", "b"]), "(a)|(b)") self.assertEqual(collect_flow._CombineRegex(["a", "b", "c"]), "(a)|(b)|(c)") self.assertEqual(collect_flow._CombineRegex(["a|b", "[^_]b", "c|d"]), "(a|b)|([^_]b)|(c|d)")
def testNewArtifactLoaded(self): """Simulate a new artifact being loaded into the store via the UI.""" cmd_artifact = """name: "TestCmdArtifact" doc: "Test command artifact for dpkg." sources: - type: "COMMAND" attributes: cmd: "/usr/bin/dpkg" args: ["--list"] labels: [ "Software" ] supported_os: [ "Linux" ] """ no_datastore_artifact = """name: "NotInDatastore" doc: "Test command artifact for dpkg." sources: - type: "COMMAND" attributes: cmd: "/usr/bin/dpkg" args: ["--list"] labels: [ "Software" ] supported_os: [ "Linux" ] """ test_registry = artifact_registry.ArtifactRegistry() test_registry.ClearRegistry() test_registry.AddDatastoreSource(rdfvalue.RDFURN("aff4:/artifact_store")) test_registry._dirty = False with utils.Stubber(artifact_registry, "REGISTRY", test_registry): collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token) with self.assertRaises(artifact_registry.ArtifactNotRegisteredError): artifact_registry.REGISTRY.GetArtifact("TestCmdArtifact") with self.assertRaises(artifact_registry.ArtifactNotRegisteredError): artifact_registry.REGISTRY.GetArtifact("NotInDatastore") # Add artifact to datastore but not registry artifact_coll = artifact_registry.ArtifactCollection( rdfvalue.RDFURN("aff4:/artifact_store")) with data_store.DB.GetMutationPool() as pool: for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml( cmd_artifact): artifact_coll.Add(artifact_val, mutation_pool=pool) # Add artifact to registry but not datastore for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml( no_datastore_artifact): artifact_registry.REGISTRY.RegisterArtifact( artifact_val, source="datastore", overwrite_if_exists=False) # This should succeeded because the artifacts will be reloaded from the # datastore. self.assertTrue(collect_flow._GetArtifactFromName("TestCmdArtifact")) # We registered this artifact with datastore source but didn't # write it into aff4. This simulates an artifact that was # uploaded in the UI then later deleted. We expect it to get # cleared when the artifacts are reloaded from the datastore. with self.assertRaises(artifact_registry.ArtifactNotRegisteredError): artifact_registry.REGISTRY.GetArtifact("NotInDatastore")
def testInterpolateArgs(self): collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token) kb = rdf_client.KnowledgeBase() kb.MergeOrAddUser(rdf_client.User(username="******")) kb.MergeOrAddUser(rdf_client.User(username="******")) collect_flow.state["knowledge_base"] = kb collect_flow.current_artifact_name = "blah" collect_flow.args = artifact_utils.ArtifactCollectorFlowArgs() test_rdf = rdf_client.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%%", r"%%users.username%%\aa"]) self.assertItemsEqual(list_args, ["test1", "test2", r"test1\aa", r"test2\aa"]) list_args = collect_flow.InterpolateList(["one"]) self.assertEqual(list_args, ["one"]) # Ignore the failure in users.desktop, report the others. collect_flow.args.ignore_interpolation_errors = True list_args = collect_flow.InterpolateList( ["%%users.desktop%%", r"%%users.username%%\aa"]) self.assertItemsEqual(list_args, [r"test1\aa", r"test2\aa"]) # Both fail. list_args = collect_flow.InterpolateList( [r"%%users.desktop%%\aa", r"%%users.sid%%\aa"]) self.assertItemsEqual(list_args, [])