def testForemanRulesAreCorrectlyRemovedWhenHuntIsStopped(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)) ]) 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) hunt.StopHunt(hunt_obj.hunt_id) rules = data_store.REL_DB.ReadAllForemanRules() self.assertEmpty(rules)
def testStoppingHuntMarksHuntFlowsForTermination(self): hunt_args = rdf_hunt_objects.HuntArguments.Standard( flow_name=compatibility.GetName(flow_test_lib.InfiniteFlow)) hunt_id, client_ids = self._CreateAndRunHunt( num_clients=5, iteration_limit=10, client_rule_set=foreman_rules.ForemanClientRuleSet(), client_rate=0, args=hunt_args) hunt.StopHunt(hunt_id) num_iterations = self._RunHunt(client_ids, iteration_limit=100) # If we did all of the planned 100 iterations, that's a clear indicator # that hunt's flows were not stopped. self.assertLess(num_iterations, 100) for client_id in client_ids: flows = data_store.REL_DB.ReadAllFlowObjects(client_id=client_id) 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, flow_obj.flow_id) self.assertFalse(req_resp)
def testRaisesWhenStoppingHuntInTheWrongState(self): hunt.StartHunt(self.hunt_id) hunt.StopHunt(self.hunt_id) self.args.state = "STOPPED" with self.assertRaises(hunt_plugin.HuntNotStoppableError): self.handler.Handle(self.args, token=self.token)
def Run(self): if data_store.RelationalDBEnabled(): replace = {} for i in range(0, 2): with test_lib.FakeTime((1 + i) * 1000): hunt_id = self.CreateHunt(description="hunt_%d" % i) if i % 2: hunt.StopHunt(hunt_id) replace[hunt_id] = "H:00000%d" % i else: replace = {} for i in range(0, 2): with test_lib.FakeTime((1 + i) * 1000): with self.CreateHunt(description="hunt_%d" % i) as hunt_obj: if i % 2: hunt_obj.Stop() replace[hunt_obj.urn.Basename()] = "H:00000%d" % i self.Check("ListHunts", args=hunt_plugin.ApiListHuntsArgs(), replace=replace) self.Check("ListHunts", args=hunt_plugin.ApiListHuntsArgs(count=1), replace=replace) self.Check("ListHunts", args=hunt_plugin.ApiListHuntsArgs(offset=1, count=1), replace=replace)
def testRaisesWhenStartingHuntInTheWrongState(self): hunt.StartHunt(self.hunt_id) hunt.StopHunt(self.hunt_id) self.args.state = "STARTED" with self.assertRaises(hunt_plugin.HuntNotStartableError): self.handler.Handle(self.args, context=self.context)
def ProcessHuntClientCrash(flow_obj, client_crash_info): """Processes client crash triggerted by a given hunt-induced flow.""" if not hunt.IsLegacyHunt(flow_obj.parent_hunt_id): def UpdateFn(hunt_obj): hunt_obj.num_crashed_clients += 1 return hunt_obj hunt_obj = data_store.REL_DB.UpdateHuntObject(flow_obj.parent_hunt_id, UpdateFn) if (hunt_obj.crash_limit and hunt_obj.num_crashed_clients >= hunt_obj.crash_limit): # Remove our rules from the forman and cancel all the started flows. # Hunt will be hard-stopped and it will be impossible to restart it. reason = ("Hunt %s reached the crashes limit of %d " "and was stopped.") % (hunt_obj.hunt_id, hunt_obj.crash_limit) hunt.StopHunt(hunt_obj.hunt_id, reason=reason) hunt_obj = hunt.StopHuntIfAverageLimitsExceeded(hunt_obj) return hunt_urn = rdfvalue.RDFURN("hunts").Add(flow_obj.parent_hunt_id) with aff4.FACTORY.Open(hunt_urn, mode="rw") as fd: # Legacy AFF4 code expects token to be set. fd.token = access_control.ACLToken(username=fd.creator) fd.RegisterCrash(client_crash_info)
def Handle(self, args, context=None): hunt_id = str(args.hunt_id) has_change = False for field_name in ["client_limit", "client_rate", "expires"]: if args.HasField(field_name): has_change = True break try: hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) if has_change: kw_args = {} if hunt_obj.hunt_state != hunt_obj.HuntState.PAUSED: raise HuntNotModifiableError( "Hunt's client limit/client rate/expiry time attributes " "can only be changed if hunt's current state is " "PAUSED") if args.HasField("client_limit"): kw_args["client_limit"] = args.client_limit if args.HasField("client_rate"): kw_args["client_rate"] = args.client_rate if args.duration is None: kw_args["duration"] = None else: kw_args["duration"] = rdfvalue.Duration(args.duration) data_store.REL_DB.UpdateHuntObject(hunt_id, **kw_args) hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id) except db.UnknownHuntError: raise HuntNotFoundError("Hunt with id %s could not be found" % args.hunt_id) if args.HasField("state"): if args.state == ApiHunt.State.STARTED: if hunt_obj.hunt_state != hunt_obj.HuntState.PAUSED: raise HuntNotStartableError( "Hunt can only be started from PAUSED state.") hunt_obj = hunt.StartHunt(hunt_obj.hunt_id) elif args.state == ApiHunt.State.STOPPED: if hunt_obj.hunt_state not in [ hunt_obj.HuntState.PAUSED, hunt_obj.HuntState.STARTED ]: raise HuntNotStoppableError( "Hunt can only be stopped from STARTED or " "PAUSED states.") hunt_obj = hunt.StopHunt(hunt_obj.hunt_id) else: raise InvalidHuntStateError( "Hunt's state can only be updated to STARTED or STOPPED") hunt_counters = data_store.REL_DB.ReadHuntCounters(hunt_obj.hunt_id) return ApiHunt().InitFromHuntObject( hunt_obj, hunt_counters=hunt_counters, with_full_summary=True)
def testRaisesWhenStoppingHuntInTheWrongState(self): if data_store.RelationalDBEnabled(): hunt.StartHunt(self.hunt_id) hunt.StopHunt(self.hunt_id) else: self.hunt_obj.Run() self.hunt_obj.Stop() self.args.state = "STOPPED" with self.assertRaises(hunt_plugin.HuntNotStoppableError): self.handler.Handle(self.args, token=self.token)
def Run(self): replace = {} for i in range(0, 2): with test_lib.FakeTime((1 + i) * 1000): hunt_id = self.CreateHunt(description="hunt_%d" % i) if i % 2: hunt.StopHunt(hunt_id) replace[hunt_id] = "H:00000%d" % i self.Check("ListHunts", args=hunt_plugin.ApiListHuntsArgs(), replace=replace) self.Check("ListHunts", args=hunt_plugin.ApiListHuntsArgs(count=1), replace=replace) self.Check("ListHunts", args=hunt_plugin.ApiListHuntsArgs(offset=1, count=1), replace=replace)