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(flow, "StartAFF4Flow", self.StartFlow): # Now setup the filters now = rdfvalue.RDFDatetime.Now() expires = now + rdfvalue.Duration("1h") foreman_obj = foreman.GetForeman(token=self.token) # Make a new rule rule = foreman_rules.ForemanRule(created=now, expires=expires, description="Test rule") # 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)) ]) # Will run Test Flow rule.actions.Append(flow_name="Test Flow", argv=rdf_protodict.Dict(foo="bar")) # Clear the rule set and add the new rule to it. rule_set = foreman_obj.Schema.RULES() rule_set.Append(rule) # Assign it to the foreman foreman_obj.Set(foreman_obj.Schema.RULES, rule_set) foreman_obj.Close() self.clients_launched = [] 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.assertEqual(len(self.clients_launched), 2) self.assertEqual(self.clients_launched[0][0], rdf_client.ClientURN(u"C.1000000000000001")) self.assertEqual(self.clients_launched[1][0], rdf_client.ClientURN(u"C.1000000000000003")) self.clients_launched = [] # 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.assertEqual(len(self.clients_launched), 0)
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 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.Duration("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(flow, "StartAFF4Flow", self.StartFlow): # Now setup the filters now = rdfvalue.RDFDatetime.Now() expires = now + rdfvalue.Duration("1h") foreman_obj = foreman.GetForeman(token=self.token) # Make a new rule rule = foreman_rules.ForemanRule(created=now, expires=expires, description="Test rule(old)") # Matches the old client one_hour_ago = base_time - rdfvalue.Duration("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())) ]) old_flow = "Test flow for old clients" # Will run Test Flow rule.actions.Append(flow_name=old_flow, argv=rdf_protodict.Dict(dict(foo="bar"))) # Clear the rule set and add the new rule to it. rule_set = foreman_obj.Schema.RULES() rule_set.Append(rule) # Make a new rule rule = foreman_rules.ForemanRule(created=now, expires=expires, description="Test rule(new)") # 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())) ]) new_flow = "Test flow for newer clients" # Will run Test Flow rule.actions.Append(flow_name=new_flow, argv=rdf_protodict.Dict(dict(foo="bar"))) rule_set.Append(rule) # Make a new rule rule = foreman_rules.ForemanRule(created=now, expires=expires, description="Test rule(eq)") # 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())) ]) eq_flow = "Test flow for LAST_BOOT_TIME" rule.actions.Append(flow_name=eq_flow, argv=rdf_protodict.Dict(dict(foo="bar"))) rule_set.Append(rule) # Assign it to the foreman foreman_obj.Set(foreman_obj.Schema.RULES, rule_set) foreman_obj.Close() self.clients_launched = [] 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.assertEqual(len(self.clients_launched), 4) self.assertEqual(self.clients_launched[0][0], rdf_client.ClientURN(u"C.1000000000000011")) self.assertEqual(self.clients_launched[0][1], new_flow) self.assertEqual(self.clients_launched[1][0], rdf_client.ClientURN(u"C.1000000000000012")) self.assertEqual(self.clients_launched[1][1], new_flow) self.assertEqual(self.clients_launched[2][0], rdf_client.ClientURN(u"C.1000000000000013")) self.assertEqual(self.clients_launched[2][1], old_flow) self.assertEqual(self.clients_launched[3][0], rdf_client.ClientURN(u"C.1000000000000014")) self.assertEqual(self.clients_launched[3][1], eq_flow)
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)