def testHuntIsStoppedWhenExpirationTimeIsReached(self): client_ids = self.SetupClients(5) duration = rdfvalue.Duration.From(1, rdfvalue.DAYS) expiry_time = rdfvalue.RDFDatetime.Now() + duration hunt_id = self._CreateHunt( client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, duration=duration, args=self.GetFileHuntArgs()) client_mock = hunt_test_lib.SampleHuntMock(failrate=-1) foreman_obj = foreman.Foreman() for client_id in client_ids: foreman_obj.AssignTasksToClient(client_id) hunt_test_lib.TestHuntHelper(client_mock, client_ids[:3]) hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, rdf_hunt_objects.Hunt.HuntState.STARTED) with test_lib.FakeTime(expiry_time - rdfvalue.Duration.From(1, rdfvalue.SECONDS)): hunt_test_lib.TestHuntHelper(client_mock, client_ids[3:4]) hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, rdf_hunt_objects.Hunt.HuntState.STARTED) with test_lib.FakeTime(expiry_time + rdfvalue.Duration.From(1, rdfvalue.SECONDS)): hunt_test_lib.TestHuntHelper(client_mock, client_ids[4:5]) hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) self.assertEqual(hunt_obj.hunt_state, rdf_hunt_objects.Hunt.HuntState.COMPLETED)
def testForemanRulesWorkCorrectlyWithStandardHunt(self): client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.OS, os=foreman_rules.ForemanOsClientRule(os_windows=True)) ]) hunt_obj = rdf_hunt_objects.Hunt(client_rule_set=client_rule_set, client_rate=0, args=self.GetFileHuntArgs()) hunt_obj.args.hunt_type = hunt_obj.args.HuntType.STANDARD data_store.REL_DB.WriteHuntObject(hunt_obj) hunt.StartHunt(hunt_obj.hunt_id) # Check matching client. client_id = self.SetupClient(0, system="Windows") foreman_obj = foreman.Foreman() foreman_obj.AssignTasksToClient(client_id) flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id) self.assertLen(flows, 1) # Check non-matching client. client_id = self.SetupClient(1, system="Linux") foreman_obj.AssignTasksToClient(client_id) flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id) self.assertEmpty(flows)
def testAssigningTasksToClientDoesNotEraseFleetspeakValidationInfo(self): client_id = self.SetupClient(0) data_store.REL_DB.WriteClientMetadata(client_id, fleetspeak_validation_info={ "foo": "bar", "12": "34" }) # Setup the rules so that AssignTasksToClient have to apply them # and update the foreman check timestamp (otherwise it will exit # doing nothing). now = rdfvalue.RDFDatetime.Now() expiration_time = now + rdfvalue.Duration.From(1, rdfvalue.HOURS) # Make a new rule rule = foreman_rules.ForemanCondition(creation_time=now, expiration_time=expiration_time, description="Test rule", hunt_id="11111111") data_store.REL_DB.WriteForemanRule(rule) foreman_obj = foreman.Foreman() foreman_obj.AssignTasksToClient(client_id) client_metadata = data_store.REL_DB.ReadClientMetadata(client_id) self.assertTrue( client_metadata.HasField("last_fleetspeak_validation_info"))
def _RunHunt(self, client_ids, client_mock=None, iteration_limit=None): foreman_obj = foreman.Foreman() for client_id in client_ids: foreman_obj.AssignTasksToClient(client_id) if client_mock is None: client_mock = hunt_test_lib.SampleHuntMock(failrate=2) return hunt_test_lib.TestHuntHelper( client_mock, client_ids, iteration_limit=iteration_limit)
def AssignTasksToClients(self, client_ids=None, worker=None): # Pretend to be the foreman now and dish out hunting jobs to all the # clients.. client_ids = client_ids or self.client_ids foreman_obj = foreman.Foreman() def Assign(): for client_id in client_ids: foreman_obj.AssignTasksToClient(client_id) if worker is None: with flow_test_lib.TestWorker(): Assign() else: Assign()
def testOperatingSystemSelection(self): """Tests that we can distinguish based on operating system.""" self.SetupClient(1, system="Windows XP") self.SetupClient(2, system="Linux") self.SetupClient(3, system="Windows 7") with utils.Stubber(hunt, "StartHuntFlowOnClient", self.StartHuntFlowOnClient): # Now setup the filters now = rdfvalue.RDFDatetime.Now() expiration_time = now + rdfvalue.DurationSeconds("1h") # Make a new rule rule = foreman_rules.ForemanCondition( creation_time=now, expiration_time=expiration_time, description="Test rule", hunt_id="111111") # Matches Windows boxes rule.client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.OS, os=foreman_rules.ForemanOsClientRule(os_windows=True)) ]) data_store.REL_DB.WriteForemanRule(rule) self.clients_started = [] foreman_obj = foreman.Foreman() foreman_obj.AssignTasksToClient(u"C.1000000000000001") foreman_obj.AssignTasksToClient(u"C.1000000000000002") foreman_obj.AssignTasksToClient(u"C.1000000000000003") # Make sure that only the windows machines ran self.assertLen(self.clients_started, 2) self.assertEqual(self.clients_started[0][1], u"C.1000000000000001") self.assertEqual(self.clients_started[1][1], u"C.1000000000000003") self.clients_started = [] # Run again - This should not fire since it did already foreman_obj.AssignTasksToClient(u"C.1000000000000001") foreman_obj.AssignTasksToClient(u"C.1000000000000002") foreman_obj.AssignTasksToClient(u"C.1000000000000003") self.assertEmpty(self.clients_started)
def testHangingClientsAreCorrectlyAccountedFor(self): client_ids = self.SetupClients(10) hunt_obj = rdf_hunt_objects.Hunt( client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=self.GetFileHuntArgs()) hunt.CreateHunt(hunt_obj) hunt_obj = hunt.StartHunt(hunt_obj.hunt_id) foreman_obj = foreman.Foreman() for client_id in client_ids: foreman_obj.AssignTasksToClient(client_id) client_mock = hunt_test_lib.SampleHuntMock(failrate=2) hunt_test_lib.TestHuntHelper(client_mock, client_ids[1:9]) hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_obj.hunt_id) self.assertEqual(hunt_counters.num_clients, 10) self.assertEqual(hunt_counters.num_successful_clients, 4) self.assertEqual(hunt_counters.num_failed_clients, 4)
def testLabelsHuntRuleMatchesCorrectClients(self): client_ids = self.SetupClients(10) self.AddClientLabel(client_ids[1], u"owner1", u"foo") self.AddClientLabel(client_ids[1], u"owner2", u"bar") self.AddClientLabel(client_ids[7], u"GRR", u"bar") self.Open("/#main=ManageHunts") self.Click("css=button[name=NewHunt]") # Select "List Processes" flow. self.Click("css=#_Processes > i.jstree-icon") self.Click("link=ListProcesses") # Click 'Next' to go to the output plugins page, hunt parameters page # and then to hunt rules page. self.Click("css=grr-new-hunt-wizard-form button.Next") self.WaitUntil(self.IsElementPresent, "css=grr-wizard-form:contains('Hunt parameters')") self.Click("css=grr-new-hunt-wizard-form button.Next") self.WaitUntil( self.IsElementPresent, "css=grr-wizard-form:contains('How to process results')") self.Click("css=grr-new-hunt-wizard-form button.Next") self.WaitUntil(self.IsElementPresent, "css=grr-wizard-form:contains('Where to run?')") # Select 'Clients With Label' rule. self.Click("css=grr-configure-rules-page button[name=Add]") self.Select("css=grr-new-hunt-wizard-form div.well select", "Label") self.Select( "css=grr-new-hunt-wizard-form div.well .form-group " ".form-group:has(label:contains('Label')):nth-last-of-type(1) " "select", "foo") self.Click("css=grr-new-hunt-wizard-form div.well .form-group " ".form-group:has(label:contains('Add label')) button") self.Select( "css=grr-new-hunt-wizard-form div.well .form-group " ".form-group:has(label:contains('Label')):nth-last-of-type(1) " "select", "bar") self.Select( "css=grr-new-hunt-wizard-form div.well .form-group " ".form-group:has(label:contains('Match mode')) select", "Match any") # Click 'Next' to go to hunt overview page. Then click 'Next' to go to # submit the hunt and wait until it's created. self.Click("css=grr-new-hunt-wizard-form button.Next") self.Click("css=grr-new-hunt-wizard-form button.Next") self.WaitUntil(self.IsElementPresent, "css=grr-wizard-form:contains('Created Hunt')") hunts_list = sorted(data_store.REL_DB.ReadHuntObjects(offset=0, count=10), key=lambda x: x.create_time) hunt = hunts_list[0] lib_hunt.StartHunt(hunt.hunt_id) foreman_obj = foreman.Foreman() for client_id in client_ids: tasks_assigned = foreman_obj.AssignTasksToClient(client_id) if client_id in [client_ids[1], client_ids[7]]: self.assertTrue(tasks_assigned) else: self.assertFalse(tasks_assigned)
def testIntegerComparisons(self): """Tests that we can use integer matching rules on the foreman.""" base_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1336480583.077736) boot_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch( 1336300000.000000) self.SetupClient(0x11, system="Windows XP", install_time=base_time) self.SetupClient(0x12, system="Windows 7", install_time=base_time) # This one was installed one week earlier. one_week_ago = base_time - rdfvalue.DurationSeconds("1w") self.SetupClient(0x13, system="Windows 7", install_time=one_week_ago) self.SetupClient(0x14, system="Windows 7", last_boot_time=boot_time) with utils.Stubber(hunt, "StartHuntFlowOnClient", self.StartHuntFlowOnClient): now = rdfvalue.RDFDatetime.Now() expiration_time = now + rdfvalue.DurationSeconds("1h") # Make a new rule rule = foreman_rules.ForemanCondition( creation_time=now, expiration_time=expiration_time, description="Test rule(old)", hunt_id="111111") # Matches the old client one_hour_ago = base_time - rdfvalue.DurationSeconds("1h") rule.client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.INTEGER, integer=foreman_rules.ForemanIntegerClientRule( field="INSTALL_TIME", operator=foreman_rules.ForemanIntegerClientRule. Operator.LESS_THAN, value=one_hour_ago.AsSecondsSinceEpoch())) ]) data_store.REL_DB.WriteForemanRule(rule) # Make a new rule rule = foreman_rules.ForemanCondition( creation_time=now, expiration_time=expiration_time, description="Test rule(new)", hunt_id="222222") # Matches the newer clients rule.client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.INTEGER, integer=foreman_rules.ForemanIntegerClientRule( field="INSTALL_TIME", operator=foreman_rules.ForemanIntegerClientRule. Operator.GREATER_THAN, value=one_hour_ago.AsSecondsSinceEpoch())) ]) data_store.REL_DB.WriteForemanRule(rule) # Make a new rule rule = foreman_rules.ForemanCondition( creation_time=now, expiration_time=expiration_time, description="Test rule(eq)", hunt_id="333333") # Note that this also tests the handling of nonexistent attributes. rule.client_rule_set = foreman_rules.ForemanClientRuleSet(rules=[ foreman_rules.ForemanClientRule( rule_type=foreman_rules.ForemanClientRule.Type.INTEGER, integer=foreman_rules.ForemanIntegerClientRule( field="LAST_BOOT_TIME", operator="EQUAL", value=boot_time.AsSecondsSinceEpoch())) ]) data_store.REL_DB.WriteForemanRule(rule) foreman_obj = foreman.Foreman() self.clients_started = [] foreman_obj.AssignTasksToClient(u"C.1000000000000011") foreman_obj.AssignTasksToClient(u"C.1000000000000012") foreman_obj.AssignTasksToClient(u"C.1000000000000013") foreman_obj.AssignTasksToClient(u"C.1000000000000014") # Make sure that the clients ran the correct flows. self.assertLen(self.clients_started, 4) self.assertEqual(self.clients_started[0][1], u"C.1000000000000011") self.assertEqual("222222", self.clients_started[0][0]) self.assertEqual(self.clients_started[1][1], u"C.1000000000000012") self.assertEqual("222222", self.clients_started[1][0]) self.assertEqual(self.clients_started[2][1], u"C.1000000000000013") self.assertEqual("111111", self.clients_started[2][0]) self.assertEqual(self.clients_started[3][1], u"C.1000000000000014") self.assertEqual("333333", self.clients_started[3][0])
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)