def testForemanRulesAreCorrectlyPropagatedWhenHuntStarts(self): client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.REGEX, regex=foreman_rules.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="HUNT")), foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.INTEGER, integer=foreman_rules.ForemanIntegerClientRule( field="CLIENT_CLOCK", operator=foreman_rules.ForemanIntegerClientRule.Operator. GREATER_THAN, value=1336650631137737)) ]) self.assertEmpty(data_store.REL_DB.ReadAllForemanRules()) hunt_obj = rdf_hunt_objects.Hunt(client_rule_set=client_rule_set) hunt_obj.args.hunt_type = hunt_obj.args.HuntType.STANDARD data_store.REL_DB.WriteHuntObject(hunt_obj) hunt_obj = hunt.StartHunt(hunt_obj.hunt_id) rules = data_store.REL_DB.ReadAllForemanRules() self.assertLen(rules, 1) rule = rules[0] self.assertEqual(rule.client_rule_set, client_rule_set) self.assertEqual(rule.hunt_id, hunt_obj.hunt_id) self.assertEqual(rule.expiration_time, hunt_obj.init_start_time + hunt_obj.duration) # Running a second time should not change the rules any more. with self.assertRaises(hunt.OnlyPausedHuntCanBeStartedError): hunt.StartHunt(hunt_obj.hunt_id) rules = data_store.REL_DB.ReadAllForemanRules() self.assertLen(rules, 1)
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 = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.REGEX, regex=foreman_rules.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="GRR")) ]) with implementation.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 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.assertEmpty(rules) now = rdfvalue.RDFDatetime.Now() expires = rdfvalue.Duration("1h").Expiry() # Add some rules. rules = [ foreman_rules.ForemanRule(created=now, expires=expires, description="Test rule1"), foreman_rules.ForemanRule(created=now, expires=expires, description="Test rule2") ] self.AddForemanRules(rules) client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.REGEX, regex=foreman_rules.ForemanRegexClientRule( field="CLIENT_NAME", attribute_regex="HUNT")), foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.INTEGER, integer=foreman_rules.ForemanIntegerClientRule( field="CLIENT_CLOCK", operator=foreman_rules.ForemanIntegerClientRule.Operator. GREATER_THAN, value=1336650631137737)) ]) hunt = implementation.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 = [ foreman_rules.ForemanRule(created=now, expires=expires, description="Test rule3"), foreman_rules.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.assertLen(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.assertLen(rules, 4) # And the hunt should report no outstanding requests any more. with hunt: self.assertFalse(hunt.GetRunner().IsHuntStarted())
def testRuleExpiration(self): with test_lib.FakeTime(1000): foreman_obj = foreman.Foreman() rules = [] rules.append( foreman_rules.ForemanCondition( creation_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1000), expiration_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1500), description="Test rule1", hunt_id="111111")) rules.append( foreman_rules.ForemanCondition( creation_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1000), expiration_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1200), description="Test rule2", hunt_id="222222")) rules.append( foreman_rules.ForemanCondition( creation_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1000), expiration_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1500), description="Test rule3", hunt_id="333333")) rules.append( foreman_rules.ForemanCondition( creation_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1000), expiration_time=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1300), description="Test rule4", hunt_id="444444")) client_id = self.SetupClient(0x21) # Clear the rule set and add the new rules to it. for rule in rules: # Add some regex that does not match the client. rule.client_rule_set = foreman_rules.ForemanClientRuleSet( rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type. REGEX, regex=foreman_rules.ForemanRegexClientRule( field="SYSTEM", attribute_regex="XXX")) ]) data_store.REL_DB.WriteForemanRule(rule) for now, num_rules in [(1000, 4), (1250, 3), (1350, 2), (1600, 0)]: with test_lib.FakeTime(now): data_store.REL_DB.WriteClientMetadata( client_id, last_foreman=rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 100)) foreman_obj.AssignTasksToClient(client_id) rules = data_store.REL_DB.ReadAllForemanRules() self.assertLen(rules, num_rules)
def testEvaluatesTheWholeAttributeToTrue(self): r = foreman_rules.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 GenerateSample(self, number=0): # Sample rule matches clients that have str(number) in their MAC return foreman_rules.ForemanRegexClientRule( field="MAC_ADDRESSES", attribute_regex=str(number))
def testRuleExpiration(self): with test_lib.FakeTime(1000): foreman_obj = foreman.GetForeman(token=self.token) hunt_id = rdfvalue.SessionID("aff4:/hunts/foremantest") rules = [] rules.append( foreman_rules.ForemanRule( created=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1000), expires=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1500), description="Test rule1")) rules.append( foreman_rules.ForemanRule( created=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1000), expires=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1200), description="Test rule2")) rules.append( foreman_rules.ForemanRule( created=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1000), expires=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1500), description="Test rule3")) rules.append( foreman_rules.ForemanRule( created=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1000), expires=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(1300), description="Test rule4", actions=[foreman_rules.ForemanRuleAction(hunt_id=hunt_id) ])) client_id = u"C.0000000000000021" fd = aff4.FACTORY.Create(client_id, aff4_grr.VFSGRRClient, token=self.token) fd.Close() # Clear the rule set and add the new rules to it. rule_set = foreman_obj.Schema.RULES() for rule in rules: # Add some regex that does not match the client. rule.client_rule_set = foreman_rules.ForemanClientRuleSet( rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type. REGEX, regex=foreman_rules.ForemanRegexClientRule( field="SYSTEM", attribute_regex="XXX")) ]) rule_set.Append(rule) foreman_obj.Set(foreman_obj.Schema.RULES, rule_set) foreman_obj.Close() fd = aff4.FACTORY.Create(client_id, aff4_grr.VFSGRRClient, token=self.token) for now, num_rules in [(1000, 4), (1250, 3), (1350, 2), (1600, 0)]: with test_lib.FakeTime(now): fd.Set(fd.Schema.LAST_FOREMAN_TIME(100)) fd.Flush() foreman_obj = foreman.GetForeman(token=self.token) foreman_obj.AssignTasksToClient(client_id) rules = foreman_obj.Get(foreman_obj.Schema.RULES) self.assertEqual(len(rules), num_rules) # Expiring rules that trigger hunts creates a notification for that hunt. with queue_manager.QueueManager(token=self.token) as manager: notifications = manager.GetNotificationsForAllShards( hunt_id.Queue()) self.assertEqual(len(notifications), 1) self.assertEqual(notifications[0].session_id, hunt_id)