def testInterpolateArgs(self): collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token) collect_flow.state.Register("knowledge_base", rdf_client.KnowledgeBase()) collect_flow.current_artifact_name = "blah" collect_flow.state.knowledge_base.MergeOrAddUser( rdf_client.KnowledgeBaseUser(username="******")) collect_flow.state.knowledge_base.MergeOrAddUser( rdf_client.KnowledgeBaseUser(username="******")) 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"])
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", rdf_client.KnowledgeBase()) collect_flow.current_artifact_name = "blah" collect_flow.state.knowledge_base.MergeOrAddUser( rdf_client.KnowledgeBaseUser(username="******")) collect_flow.state.knowledge_base.MergeOrAddUser( rdf_client.KnowledgeBaseUser(username="******")) collector = artifact_registry.ArtifactSource( type=artifact_registry.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 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 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._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 for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml( cmd_artifact): with aff4.FACTORY.Create( "aff4:/artifact_store", aff4_type=collects.RDFValueCollection, token=self.token, mode="rw") as artifact_coll: artifact_coll.Add(artifact_val) # 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 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 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, [])