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 = 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_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 CreateHunt(self, flow_runner_args=None, flow_args=None, client_rule_set=None, original_object=None, client_rate=0, token=None, **kwargs): # Only initialize default flow_args value if default flow_runner_args value # is to be used. if not flow_runner_args: flow_args = (flow_args or transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS))) flow_runner_args = (flow_runner_args or rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__)) client_rule_set = (client_rule_set or self._CreateForemanClientRuleSet()) return implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=flow_runner_args, flow_args=flow_args, client_rule_set=client_rule_set, client_rate=client_rate, original_object=original_object, token=token or self.token, **kwargs)
def testNoKillNotificationsScheduledForHunts(self): worker_obj = self._TestWorker() initial_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100) try: with test_lib.FakeTime(initial_time.AsSecondsSinceEpoch()): with implementation.StartHunt( hunt_name=WorkerStuckableHunt.__name__, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() implementation.GRRHunt.StartClients(hunt.session_id, [self.client_id]) # Process all messages while worker_obj.RunOnce(): pass # Wait until worker thread starts processing the flow. WorkerStuckableHunt.WaitUntilWorkerStartsProcessing() # Assert that there are no stuck notifications in the worker's queue. with queue_manager.QueueManager(token=self.token) as manager: for queue in worker_obj.queues: notifications = manager.GetNotifications(queue) for n in notifications: self.assertFalse(n.in_progress) finally: # Release the semaphore so that worker thread unblocks and finishes # processing the flow. WorkerStuckableHunt.LetWorkerFinishProcessing() worker_obj.thread_pool.Join()
def testOverviewIsShownForNestedHuntFlows(self): if data_store.RelationalDBFlowsEnabled(): # TODO(amoser): Hunts don't spawn relational flows yet. return with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=gui_test_lib.RecursiveTestFlow.__name__), client_rate=0, token=self.token) as hunt: hunt.Run() self.AssignTasksToClients(client_ids=[self.client_id]) self.RunHunt(client_ids=[self.client_id]) self.Open("/#/clients/%s" % self.client_id) self.Click("css=a[grrtarget='client.flows']") # There should be a RecursiveTestFlow in the list. Expand nested flows. self.Click("css=tr:contains('RecursiveTestFlow') span.tree_branch") # Click on a nested flow. self.Click("css=tr:contains('RecursiveTestFlow'):nth(2)") # Nested flow should have Depth argument set to 1. self.WaitUntil(self.IsElementPresent, "css=td:contains('Depth') ~ td:nth(0):contains('1')") # Check that flow id of this flow has forward slash - i.e. consists of # 2 components. self.WaitUntil(self.IsTextPresent, "Flow ID") flow_id = self.GetText("css=dt:contains('Flow ID') ~ dd:nth(0)") self.assertIn("/", flow_id)
def StartInterrogationHunt(self): """Starts an interrogation hunt on all available clients.""" flow_name = compatibility.GetName(flows_discovery.Interrogate) flow_args = flows_discovery.InterrogateArgs(lightweight=False) description = "Interrogate run by cron to keep host info fresh." if data_store.RelationalDBEnabled(): hunt_id = hunt.CreateAndStartHunt( flow_name, flow_args, self.token.username, client_limit=0, client_rate=50, crash_limit=config.CONFIG["Cron.interrogate_crash_limit"], description=description, duration=rdfvalue.DurationSeconds("1w"), output_plugins=self.GetOutputPlugins()) self.Log("Started hunt %s.", hunt_id) else: with hunts_implementation.StartHunt( hunt_name=hunts_standard.GenericHunt.__name__, client_limit=0, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=flow_name), flow_args=flow_args, output_plugins=self.GetOutputPlugins(), crash_limit=config.CONFIG["Cron.interrogate_crash_limit"], client_rate=50, expiry_time=rdfvalue.DurationSeconds("1w"), description=description, token=self.token) as hunt_obj: hunt_obj.GetRunner().Start() self.Log("Started hunt %s.", hunt_obj.urn)
def setUp(self): super(ApiGetHuntFileHandlerTest, self).setUp() self.handler = hunt_plugin.ApiGetHuntFileHandler() self.file_path = os.path.join(self.base_path, "test.plist") self.hunt = implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=file_finder.FileFinder.__name__), flow_args=rdf_file_finder.FileFinderArgs( paths=[self.file_path], action=rdf_file_finder.FileFinderAction( action_type="DOWNLOAD"), ), client_rate=0, token=self.token) self.hunt.Run() self.aff4_file_path = "fs/os/%s" % self.file_path self.client_id = self.SetupClient(0) self.AssignTasksToClients(client_ids=[self.client_id]) action_mock = action_mocks.FileFinderClientMock() hunt_test_lib.TestHuntHelper(action_mock, [self.client_id], token=self.token)
def testStartClients(self): client_id = test_lib.TEST_CLIENT_ID with implementation.StartHunt(hunt_name=standard.SampleHunt.__name__, client_rate=0, token=self.token) as hunt: hunt.GetRunner().Start() flows = list( aff4.FACTORY.Open(client_id.Add("flows"), token=self.token).ListChildren()) self.assertEqual(flows, []) implementation.GRRHunt.StartClients(hunt.session_id, [client_id]) hunt_test_lib.TestHuntHelper(None, [client_id], False, self.token) flows = list( aff4.FACTORY.Open(client_id.Add("flows"), token=self.token).ListChildren()) # One flow should have been started. self.assertLen(flows, 1) self.assertIn(hunt.session_id.Basename(), str(flows[0]))
def testCreatorPropagation(self): self.CreateAdminUser("adminuser") admin_token = access_control.ACLToken( username="******", reason="testing") # Start a flow that requires admin privileges in the hunt. The # parameters are not valid so the flow will error out but it's # enough to check if the flow was actually run (i.e., it passed # the label test). with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=administrative.UpdateClient.__name__), flow_args=administrative.UpdateClientArgs(), client_rule_set=self._CreateForemanClientRuleSet(), client_rate=0, token=admin_token) as hunt: hunt.Run() self.CreateUser("nonadmin") nonadmin_token = access_control.ACLToken( username="******", reason="testing") self.AssignTasksToClients() client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, nonadmin_token) errors = list(hunt.GetClientsErrors()) # Make sure there are errors... self.assertTrue(errors) # but they are not UnauthorizedAccess. for e in errors: self.assertNotIn("UnauthorizedAccess", e.backtrace)
def CreateSampleHunt(self): """Creats SampleHunt, writes it to the data store and returns it's id.""" with implementation.StartHunt( hunt_name=standard.SampleHunt.__name__, token=self.token.SetUID()) as hunt: return hunt.session_id
def SetUpCrashedFlowInHunt(self): client_ids = self.SetupClients(10) client_mocks = dict([ (client_id, flow_test_lib.CrashClientMock(client_id, self.token)) for client_id in client_ids ]) client_rule_set = self._CreateForemanClientRuleSet() # Make this not match anything. client_rule_set.rules[0].regex.attribute_regex = "" with implementation.StartHunt(hunt_name=standard.SampleHunt.__name__, client_rule_set=client_rule_set, client_rate=0, token=self.token) as hunt: hunt.Run() foreman_obj = foreman.GetForeman(token=self.token) for client_id in client_ids: self.assertTrue( foreman_obj.AssignTasksToClient(client_id.Basename())) hunt_test_lib.TestHuntHelperWithMultipleMocks(client_mocks, False, self.token) return client_ids
def CreateSampleHunt(self, path=None, stopped=False, output_plugins=None, client_limit=0, client_count=10, token=None): token = token or self.token self.client_ids = self.SetupClients(client_count) with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs(pathspec=rdf_paths.PathSpec( path=path or "/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS, )), client_rule_set=self._CreateForemanClientRuleSet(), output_plugins=output_plugins or [], client_rate=0, client_limit=client_limit, token=token) as hunt: if not stopped: hunt.Run() foreman_obj = foreman.GetForeman(token=token) for client_id in self.client_ids: foreman_obj.AssignTasksToClient(client_id.Basename()) self.hunt_urn = hunt.urn return aff4.FACTORY.Open(hunt.urn, mode="rw", token=token, age=aff4.ALL_TIMES)
def testResultCounting(self): path = os.path.join(self.base_path, "hello*") num_files = len(glob.glob(path)) self.assertGreater(num_files, 0) hunt = implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=file_finder.FileFinder.__name__), flow_args=rdf_file_finder.FileFinderArgs( paths=[path], action=rdf_file_finder.FileFinderAction(action_type="STAT"), ), client_rate=0, token=self.token) hunt.Run() client_ids = self.SetupClients(5) self.AssignTasksToClients(client_ids=client_ids) action_mock = action_mocks.FileFinderClientMock() hunt_test_lib.TestHuntHelper(action_mock, client_ids, token=self.token) hunt = aff4.FACTORY.Open(hunt.urn, token=self.token) self.assertEqual(hunt.context.clients_with_results_count, 5) self.assertEqual(hunt.context.results_count, 5 * num_files)
def CreateGenericHuntWithCollection(self, values=None): self.client_ids = self.SetupClients(10) CreateFileVersion(self.client_ids[0], "fs/os/c/bin/bash", token=self.token) if values is None: values = [ rdfvalue.RDFURN("aff4:/sample/1"), rdfvalue.RDFURN("aff4:/%s/fs/os/c/bin/bash" % self.client_ids[0].Basename()), rdfvalue.RDFURN("aff4:/sample/3") ] with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, client_rule_set=self._CreateForemanClientRuleSet(), output_plugins=[], token=self.token) as hunt: runner = hunt.GetRunner() runner.Start() collection = hunt.ResultCollection() with data_store.DB.GetMutationPool() as pool: for value in values: collection.Add(rdf_flows.GrrMessage( payload=value, source=self.client_ids[0]), mutation_pool=pool) return hunt.urn
def testAlertEmailIsSentWhenClientKilledDuringHunt(self): """Test that client killed messages are handled correctly for hunts.""" client_id = test_lib.TEST_CLIENT_ID self.email_messages = [] def SendEmail(address, sender, title, message, **_): self.email_messages.append( dict(address=address, sender=sender, title=title, message=message)) with hunts_implementation.StartHunt( hunt_name=hunts_standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=flow_test_lib.FlowWithOneClientRequest.__name__), client_rate=0, crash_alert_email="*****@*****.**", token=self.token) as hunt: hunt.Run() hunt.StartClients(hunt.session_id, client_id) with utils.Stubber(email_alerts.EMAIL_ALERTER, "SendEmail", SendEmail): client = flow_test_lib.CrashClientMock(client_id, self.token) hunt_test_lib.TestHuntHelper(client, [client_id], token=self.token, check_flow_errors=False) self.assertEqual(len(self.email_messages), 2) self.assertListEqual([ self.email_messages[0]["address"], self.email_messages[1]["address"] ], ["*****@*****.**", config.CONFIG["Monitoring.alert_email"]])
def testHuntModificationWorksCorrectly(self): """This tests running the hunt on some clients.""" with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS),), client_rule_set=self._CreateForemanClientRuleSet(), client_limit=1, client_rate=0, token=self.token) as hunt: hunt.Run() # Forget about hunt object, we'll use AFF4 for everything. hunt_session_id = hunt.session_id hunt = None # Pretend to be the foreman now and dish out hunting jobs to all the # client.. self.AssignTasksToClients() # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, self.token) # Re-open the hunt to get fresh data. hunt_obj = aff4.FACTORY.Open( hunt_session_id, age=aff4.ALL_TIMES, token=self.token) # There should be only one client, due to the limit started, _, _ = hunt_obj.GetClientsCounts() self.assertEqual(started, 1) # Check the hunt is paused. self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "PAUSED") with aff4.FACTORY.Open( hunt_session_id, mode="rw", token=self.token) as hunt_obj: runner = hunt_obj.GetRunner() runner.runner_args.client_limit = 10 runner.Start() # Pretend to be the foreman now and dish out hunting jobs to all the # clients. self.AssignTasksToClients() hunt_test_lib.TestHuntHelper(client_mock, self.client_ids, False, self.token) hunt_obj = aff4.FACTORY.Open( hunt_session_id, age=aff4.ALL_TIMES, token=self.token) # There should be only one client, due to the limit started, _, _ = hunt_obj.GetClientsCounts() self.assertEqual(started, 10)
def testResourceUsageStats(self): client_ids = self.SetupClients(10) with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS, )), client_rule_set=self._CreateForemanClientRuleSet(), output_plugins=[], client_rate=0, token=self.token) as hunt: hunt.Run() self.AssignTasksToClients(client_ids=client_ids) client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper(client_mock, client_ids, False, self.token) hunt = aff4.FACTORY.Open( hunt.urn, aff4_type=standard.GenericHunt, token=self.token) # This is called once for each state method. Each flow above runs the # Start and the StoreResults methods. usage_stats = hunt.context.usage_stats self.assertEqual(usage_stats.user_cpu_stats.num, 10) self.assertTrue(math.fabs(usage_stats.user_cpu_stats.mean - 5.5) < 1e-7) self.assertTrue( math.fabs(usage_stats.user_cpu_stats.std - 2.8722813) < 1e-7) self.assertEqual(usage_stats.system_cpu_stats.num, 10) self.assertTrue(math.fabs(usage_stats.system_cpu_stats.mean - 11) < 1e-7) self.assertTrue( math.fabs(usage_stats.system_cpu_stats.std - 5.7445626) < 1e-7) self.assertEqual(usage_stats.network_bytes_sent_stats.num, 10) self.assertTrue( math.fabs(usage_stats.network_bytes_sent_stats.mean - 16.5) < 1e-7) self.assertTrue( math.fabs(usage_stats.network_bytes_sent_stats.std - 8.61684396) < 1e-7) # NOTE: Not checking histograms here. RunningStatsTest tests that mean, # standard deviation and histograms are calculated correctly. Therefore # if mean/stdev values are correct histograms should be ok as well. self.assertLen(usage_stats.worst_performers, 10) prev = usage_stats.worst_performers[0] for p in usage_stats.worst_performers[1:]: self.assertTrue( prev.cpu_usage.user_cpu_time + prev.cpu_usage.system_cpu_time > p.cpu_usage.user_cpu_time + p.cpu_usage.system_cpu_time) prev = p
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.assertEmpty(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) # 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.assertLen(rules, 1) rule = rules[0] self.assertEqual(rule.client_rule_set, client_rule_set) self.assertLen(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.assertLen(rules, 1)
def CreateSampleHunt(self): """Creats SampleHunt, writes it to the data store and returns it's id.""" if data_store.RelationalDBEnabled(): return self.CreateHunt() else: with implementation.StartHunt( hunt_name=standard.SampleHunt.__name__, token=self.token.SetUID()) as hunt: return hunt.session_id.Basename()
def Start(self): """Create the hunt and run it.""" with implementation.StartHunt(runner_args=self.args.hunt_runner_args, args=self.args.hunt_args, token=self.token) as hunt: hunt.Run() self.Log("User %s created a new %s hunt (%s)", self.token.username, hunt.args.flow_runner_args.flow_name, hunt.urn)
def CreateSampleHunt(self, token=None): with implementation.StartHunt( hunt_name=standard.SampleHunt.__name__, client_rate=100, filename="TestFilename", client_rule_set=self._CreateForemanClientRuleSet(), token=token or self.token) as hunt: return hunt.session_id
def _StartHunt(self): with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=flow_test_lib.FlowWithOneNestedFlow.__name__), client_rate=0, token=self.token) as hunt: hunt.Run() self.AssignTasksToClients(client_ids=[self.client_urn]) self.RunHunt(client_ids=[self.client_urn], failrate=2)
def testStoppingHuntMarksFlowsForTerminationAndCleansQueues(self): with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name="InfiniteFlow"), client_rule_set=self._CreateForemanClientRuleSet(), client_rate=0, token=self.token) as hunt: hunt.Run() self.AssignTasksToClients() # Run long enough for InfiniteFlows to start. self.RunHunt( iteration_limit=len(self.client_ids) * 2, user_cpu_time=0, system_cpu_time=0, network_bytes_sent=0) self.StopHunt(hunt.urn) # All flows states should be destroyed by now. # If something is wrong with the GenericHunt.Stop implementation, # this will run forever. self.RunHunt(user_cpu_time=0, system_cpu_time=0, network_bytes_sent=0) if data_store.RelationalDBFlowsEnabled(): for client_id in self.client_ids: flows = data_store.REL_DB.ReadAllFlowObjects(client_id.Basename()) self.assertLen(flows, 1) flow_obj = flows[0] self.assertEqual(flow_obj.flow_state, flow_obj.FlowState.ERROR) self.assertEqual(flow_obj.error_message, "Parent hunt stopped.") req_resp = data_store.REL_DB.ReadAllFlowRequestsAndResponses( client_id.Basename(), flow_obj.flow_id) self.assertFalse(req_resp) else: for client_id in self.client_ids: flows_root = aff4.FACTORY.Open(client_id.Add("flows"), token=self.token) flows_list = list(flows_root.ListChildren()) # Only one flow (issued by the hunt) is expected. self.assertLen(flows_list, 1) # Check that flow's queues are deleted. with queue_manager.QueueManager(token=self.token) as manager: req_resp = list(manager.FetchRequestsAndResponses(flows_list[0])) self.assertFalse(req_resp) flow_obj = aff4.FACTORY.Open( flows_list[0], aff4_type=aff4_flows.InfiniteFlow, token=self.token) self.assertEqual( flow_obj.Get(flow_obj.Schema.PENDING_TERMINATION).reason, "Parent hunt stopped.")
def setUp(self): super(CleanHuntsFlowTest, self).setUp() self.hunts_urns = [] with test_lib.FakeTime(40): for i in range(self.NUM_HUNTS): hunt = implementation.StartHunt( hunt_name=standard.SampleHunt.__name__, expiry_time=rdfvalue.Duration("1m") * i, token=self.token) hunt.Run() self.hunts_urns.append(hunt.urn)
def Start(self): """Create the hunt, in the paused state.""" # Anyone can create the hunt but it will be created in the paused # state. Permissions are required to actually start it. with implementation.StartHunt(runner_args=self.args.hunt_runner_args, args=self.args.hunt_args, token=self.token) as hunt: # Nothing really to do here - hunts are always created in the paused # state. self.Log("User %s created a new %s hunt (%s)", self.token.username, hunt.args.flow_runner_args.flow_name, hunt.urn)
def CreateHunt(self, flow_runner_args=None, flow_args=None, client_rule_set=None, original_object=None, client_rate=0, duration=None, token=None, **kwargs): # Only initialize default flow_args value if default flow_runner_args value # is to be used. if not flow_runner_args: flow_args = ( flow_args or transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS))) flow_runner_args = ( flow_runner_args or rdf_flow_runner.FlowRunnerArgs(flow_name=transfer.GetFile.__name__)) client_rule_set = (client_rule_set or self._CreateForemanClientRuleSet()) if data_store.RelationalDBEnabled(): token = token or self.token hunt_args = rdf_hunt_objects.HuntArguments( hunt_type=rdf_hunt_objects.HuntArguments.HuntType.STANDARD, standard=rdf_hunt_objects.HuntArgumentsStandard( flow_name=flow_runner_args.flow_name, flow_args=flow_args)) hunt_obj = rdf_hunt_objects.Hunt( creator=token.username, client_rule_set=client_rule_set, original_object=original_object, client_rate=client_rate, duration=duration, args=hunt_args, **kwargs) hunt.CreateHunt(hunt_obj) return hunt_obj.hunt_id return implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=flow_runner_args, flow_args=flow_args, client_rule_set=client_rule_set, client_rate=client_rate, original_object=original_object, token=token or self.token, **kwargs)
def _CreateHunt(self, description): output_plugins = [ rdf_output_plugin.OutputPluginDescriptor(plugin_name="TestOutputPlugin") ] with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), output_plugins=output_plugins, description=description, client_rate=0, token=self.token) as hunt: return hunt
def Run(self): action = self.job.args.hunt_cron_action token = access_control.ACLToken(username="******") hunt_args = rdf_hunts.GenericHuntArgs( flow_args=action.flow_args, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=action.flow_name)) with implementation.StartHunt(hunt_name=GenericHunt.__name__, args=hunt_args, runner_args=action.hunt_runner_args, token=token) as hunt: hunt.Run()
def testRendersRequestedHuntAppoval(self): with implementation.StartHunt(hunt_name=standard.SampleHunt.__name__, token=self.token) as hunt: pass self.RequestHuntApproval(hunt.urn.Basename(), reason=self.token.reason, approver=u"approver", requestor=self.token.username) args = user_plugin.ApiListHuntApprovalsArgs() result = self.handler.Handle(args, token=self.token) self.assertLen(result.items, 1)
def testHuntExpiration(self): """This tests that hunts with a client limit terminate correctly.""" with test_lib.FakeTime(1000): with implementation.StartHunt( hunt_name=standard.GenericHunt.__name__, flow_runner_args=rdf_flow_runner.FlowRunnerArgs( flow_name=transfer.GetFile.__name__), flow_args=transfer.GetFileArgs( pathspec=rdf_paths.PathSpec( path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS)), client_rule_set=self._CreateForemanClientRuleSet(), client_limit=5, expiry_time=rdfvalue.Duration("1000s"), token=self.token) as hunt: hunt.Run() # Pretend to be the foreman now and dish out hunting jobs to all the # clients (Note we have 10 clients here). self.AssignTasksToClients() hunt_obj = aff4.FACTORY.Open( hunt.session_id, age=aff4.ALL_TIMES, token=self.token) self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "STARTED") # Now advance the time such that the hunt expires. time.time = lambda: 5000 # Run the hunt. client_mock = hunt_test_lib.SampleHuntMock() hunt_test_lib.TestHuntHelper( client_mock, self.client_ids, check_flow_errors=False, token=self.token) # No client should be processed since the hunt is expired. started, finished, errors = hunt_obj.GetClientsCounts() self.assertEqual(started, 0) self.assertEqual(finished, 0) self.assertEqual(errors, 0) hunt_obj = aff4.FACTORY.Open( hunt.session_id, age=aff4.ALL_TIMES, token=self.token) # Hunts are automatically stopped when they expire. self.assertEqual(hunt_obj.Get(hunt_obj.Schema.STATE), "COMPLETED")
def testPausingAndRestartingDoesNotStartHuntTwiceOnTheSameClient(self): """This tests if the hunt completes when some clients hang or raise.""" 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_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(failrate=2) 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)