def testEvaluatesTheWholeAttributeToTrue(self): r = rdf_foreman.ForemanRegexClientRule(field="SYSTEM", attribute_regex="^Linux$") client_id = self.SetupClient(0, system="Linux") self.assertTrue( r.Evaluate(aff4.FACTORY.Open(client_id, token=self.token)))
def testUnsetFieldRaises(self): client = self.SetupTestClientObject(0, system="Linux") info = data_store.REL_DB.ReadClientFullInfo(client.client_id) r = rdf_foreman.ForemanRegexClientRule(attribute_regex="foo") with self.assertRaises(ValueError): r.Evaluate(info)
def _CreateForemanClientRuleSet(self): return rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ])
def testEvaluatesTheWholeAttributeToTrue(self): r = rdf_foreman.ForemanRegexClientRule(field="SYSTEM", attribute_regex="^Linux$") client = self.SetupTestClientObject(0, system="Linux") info = data_store.REL_DB.ReadClientFullInfo(client.client_id) self.assertTrue(r.Evaluate(info))
def testUnsetFieldRaises(self): client_id = self.SetupClient(0, system="Linux") client = aff4.FACTORY.Open(client_id, token=self.token) r = rdf_foreman.ForemanRegexClientRule(attribute_regex="foo") with self.assertRaises(ValueError): r.Evaluate(client)
def testRuleAdding(self): foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) rules = foreman.Get(foreman.Schema.RULES) # Make sure there are no rules yet in the foreman. self.assertEqual(len(rules), 0) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="HUNT")), rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.INTEGER, integer=rdf_foreman.ForemanIntegerClientRule( field="CLIENT_CLOCK", operator=rdf_foreman.ForemanIntegerClientRule.Operator. GREATER_THAN, value=1336650631137737)) ]) hunt = implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) # Push the rules to the foreman. with hunt: hunt.GetRunner().Start() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) rules = foreman.Get(foreman.Schema.RULES) # Make sure they were written correctly. self.assertEqual(len(rules), 1) rule = rules[0] self.assertEqual(rule.client_rule_set, client_rule_set) self.assertEqual(len(rule.actions), 1) self.assertEqual(rule.actions[0].hunt_name, "SampleHunt") # Running a second time should not change the rules any more. with hunt: hunt.GetRunner().Start() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) rules = foreman.Get(foreman.Schema.RULES) # Still just one rule. self.assertEqual(len(rules), 1)
def testEvaluatesAttributesSubstringToTrue(self): r = rdf_foreman.ForemanRegexClientRule(field="SYSTEM", attribute_regex="inu") client_id = self.SetupClient(0, system="Linux") # The system contains the substring inu self.assertTrue( r.Evaluate(aff4.FACTORY.Open(client_id, token=self.token)))
def testEvaluatesNonSubstringToFalse(self): r = rdf_foreman.ForemanRegexClientRule(field="SYSTEM", attribute_regex="foo") client = self.SetupTestClientObject(0, system="Linux") info = data_store.REL_DB.ReadClientFullInfo(client.client_id) # The system doesn't contain foo self.assertFalse(r.Evaluate(info))
def testEvaluatesNonSubstringToFalse(self): r = rdf_foreman.ForemanRegexClientRule(field="SYSTEM", attribute_regex="foo") client_id = self.SetupClient(0, system="Linux") # The system doesn't contain foo self.assertFalse( r.Evaluate(aff4.FACTORY.Open(client_id, token=self.token)))
def testEvaluatesAttributesSubstringToTrue(self): r = rdf_foreman.ForemanRegexClientRule(field="SYSTEM", attribute_regex="inu") client = self.SetupTestClientObject(0, system="Linux") info = data_store.REL_DB.ReadClientFullInfo(client.client_id) # The system contains the substring inu self.assertTrue(r.Evaluate(info))
def testLabels(self): r = rdf_foreman.ForemanRegexClientRule(field="CLIENT_LABELS", attribute_regex="ell") client_id = self.SetupClient(0, system="Linux") client = aff4.FACTORY.Open(client_id, mode="rw", token=self.token) self.assertFalse(r.Evaluate(client)) client.SetLabels(["hello", "world"], owner="GRR") self.assertTrue(r.Evaluate(client))
def testEvaluation(self): now = rdfvalue.RDFDatetime().Now() client = self.SetupTestClientObject(0, last_boot_time=now) info = data_store.REL_DB.ReadClientFullInfo(client.client_id) for f in rdf_foreman.ForemanRegexClientRule.ForemanStringField.enum_dict: if f == "UNSET": continue r = rdf_foreman.ForemanRegexClientRule(field=f, attribute_regex=".") r.Evaluate(info)
def testLabels(self): client = self.SetupTestClientObject(0, system="Linux") data_store.REL_DB.AddClientLabels(client.client_id, "GRR", ["hello", "world"]) info = data_store.REL_DB.ReadClientFullInfo(client.client_id) # Match a system label. r = rdf_foreman.ForemanRegexClientRule(field="CLIENT_LABELS", attribute_regex="label1") self.assertTrue(r.Evaluate(info)) # Match a user label. r = rdf_foreman.ForemanRegexClientRule(field="CLIENT_LABELS", attribute_regex="ell") self.assertTrue(r.Evaluate(info)) # This rule doesn't match any label. r = rdf_foreman.ForemanRegexClientRule( field="CLIENT_LABELS", attribute_regex="NonExistentLabel") self.assertFalse(r.Evaluate(info))
def testPausingAndRestartingDoesNotStartHuntTwiceOnTheSameClient(self): """This tests if the hunt completes when some clients hang or raise.""" client_ids = self.SetupClients(10) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() hunt_id = hunt.urn foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: num_tasks = foreman.AssignTasksToClient(client_id.Basename()) self.assertEqual(num_tasks, 1) client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token) # Pausing and running hunt: this leads to the fresh rules being written # to Foreman.RULES. with aff4.FACTORY.Open(hunt_id, mode="rw", token=self.token) as hunt: runner = hunt.GetRunner() runner.Pause() runner.Start() # Recreating the foreman so that it updates list of rules. foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: num_tasks = foreman.AssignTasksToClient(client_id.Basename()) # No tasks should be assigned as this hunt ran on all the clients # before. self.assertEqual(num_tasks, 0)
def testInvalidRules(self): """Tests the behavior when the field is left UNSET in a rule.""" client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="UNSET", attribute_regex="HUNT")) ]) with implementation.GRRHunt.StartHunt( hunt_name=BrokenSampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: runner = hunt.GetRunner() self.assertRaises(ValueError, runner.Start)
def testBrokenHunt(self): """This tests the behavior when a hunt raises an exception.""" # Set up 10 clients. client_ids = self.SetupClients(10) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=BrokenSampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: foreman.AssignTasksToClient(client_id.Basename()) # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token) hunt_obj = aff4.FACTORY.Open(hunt.session_id, mode="rw", age=aff4.ALL_TIMES, token=self.token) started, finished, errors = hunt_obj.GetClientsCounts() self.assertEqual(started, 10) # There should be errors for the five clients where the hunt raised. self.assertEqual(errors, 5) # All of the clients that have the file should still finish eventually. self.assertEqual(finished, 5)
def testClientLimit(self): """This tests that we can limit hunts to a number of clients.""" # Set up 10 clients. client_ids = self.SetupClients(10) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_limit=5, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: hunt.Run() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: foreman.AssignTasksToClient(client_id.Basename()) # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token) hunt_obj = aff4.FACTORY.Open(hunt.urn, mode="rw", age=aff4.ALL_TIMES, token=self.token) started, finished, _ = hunt_obj.GetClientsCounts() # We limited here to 5 clients. self.assertEqual(started, 5) self.assertEqual(finished, 5)
def testHangingClients(self): """This tests if the hunt completes when some clients hang or raise.""" # Set up 10 clients. client_ids = self.SetupClients(10) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: foreman.AssignTasksToClient(client_id.Basename()) client_mock = hunt_test_lib.SampleHuntMock() # Just pass 8 clients to run, the other two went offline. hunt_test_lib.TestHuntHelper(client_mock, client_ids[1:9], False, self.token) hunt_obj = aff4.FACTORY.Open(hunt.session_id, mode="rw", age=aff4.ALL_TIMES, token=self.token) started, finished, _ = hunt_obj.GetClientsCounts() # We started the hunt on 10 clients. self.assertEqual(started, 10) # But only 8 should have finished. self.assertEqual(finished, 8)
def testProcessing(self): """This tests running the hunt on some clients.""" # Set up 10 clients. client_ids = self.SetupClients(10) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: foreman.AssignTasksToClient(client_id.Basename()) # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token) hunt_obj = aff4.FACTORY.Open(hunt.session_id, mode="r", age=aff4.ALL_TIMES, aff4_type=standard.SampleHunt, token=self.token) started, finished, _ = hunt_obj.GetClientsCounts() self.assertEqual(started, 10) self.assertEqual(finished, 10)
def _RunRateLimitedHunt(self, client_ids, start_time): client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt(hunt_name=DummyHunt.__name__, client_rule_set=client_rule_set, client_rate=1, token=self.token) as hunt: hunt.Run() # Pretend to be the foreman now and dish out hunting jobs to all the # clients.. foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) for client_id in client_ids: foreman.AssignTasksToClient(client_id.Basename()) self.assertEqual(len(DummyHunt.client_ids), 0) # Run the hunt. worker_mock = worker_test_lib.MockWorker(check_flow_errors=True, queues=queues.HUNTS, token=self.token) # One client is scheduled in the first minute. with test_lib.FakeTime(start_time + 2): worker_mock.Simulate() self.assertEqual(len(DummyHunt.client_ids), 1) # No further clients will be scheduled until the end of the first minute. with test_lib.FakeTime(start_time + 59): worker_mock.Simulate() self.assertEqual(len(DummyHunt.client_ids), 1) return worker_mock, hunt.urn
def testCallback(self, client_limit=None): """Checks that the foreman uses the callback specified in the action.""" client_urn = self.SetupClient(0) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_limit=client_limit, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() # Create a client that matches our regex. with aff4.FACTORY.Open(client_urn, mode="rw", token=self.token) as client: info = client.Schema.CLIENT_INFO() info.client_name = "GRR Monitor" client.Set(client.Schema.CLIENT_INFO, info) foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) with utils.Stubber(standard.SampleHunt, "StartClients", self.Callback): self.called = [] client_id = client_urn.Basename() foreman.AssignTasksToClient(client_id) self.assertEqual(len(self.called), 1) self.assertEqual(self.called[0][1], [client_id])
def testStopping(self): """Tests if we can stop a hunt.""" foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) rules = foreman.Get(foreman.Schema.RULES) # Make sure there are no rules yet. self.assertEqual(len(rules), 0) now = rdfvalue.RDFDatetime.Now() expires = rdfvalue.Duration("1h").Expiry() # Add some rules. rules = [ rdf_foreman.ForemanRule(created=now, expires=expires, description="Test rule1"), rdf_foreman.ForemanRule(created=now, expires=expires, description="Test rule2") ] self.AddForemanRules(rules) client_rule_set = rdf_foreman.ForemanClientRuleSet(rules=[ rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.REGEX, regex=rdf_foreman.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="HUNT")), rdf_foreman.ForemanClientRule( rule_type=rdf_foreman.ForemanClientRule.Type.INTEGER, integer=rdf_foreman.ForemanIntegerClientRule( field="CLIENT_CLOCK", operator=rdf_foreman.ForemanIntegerClientRule.Operator. GREATER_THAN, value=1336650631137737)) ]) hunt = implementation.GRRHunt.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) with hunt: runner = hunt.GetRunner() runner.Start() # Add some more rules. rules = [ rdf_foreman.ForemanRule(created=now, expires=expires, description="Test rule3"), rdf_foreman.ForemanRule(created=now, expires=expires, description="Test rule4") ] self.AddForemanRules(rules) foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) rules = foreman.Get(foreman.Schema.RULES) self.assertEqual(len(rules), 5) # It should be running. self.assertTrue(runner.IsHuntStarted()) # Now we stop the hunt. hunt.Stop() foreman = aff4.FACTORY.Open("aff4:/foreman", mode="rw", token=self.token) rules = foreman.Get(foreman.Schema.RULES) # The rule for this hunt should be deleted but the rest should be there. self.assertEqual(len(rules), 4) # And the hunt should report no outstanding requests any more. with hunt: self.assertFalse(hunt.GetRunner().IsHuntStarted())
def GenerateSample(self, number=0): # Sample rule matches clients that have str(number) in their MAC return rdf_foreman.ForemanRegexClientRule(field="MAC_ADDRESSES", attribute_regex=str(number))