def WriteFlowObject(*args, **kwargs): with mock.patch.object(data_store.REL_DB.delegate, "WriteFlowObject", original): try: hunt.StartHuntFlowOnClient(client_id, hunt_id) except Exception as e: raise AssertionError(e) return data_store.REL_DB.WriteFlowObject(*args, **kwargs)
def _RunAction(self, rule, client_id): """Run all the actions specified in the rule. Args: rule: Rule which actions are to be executed. client_id: Id of a client where rule's actions are to be executed. Returns: Number of actions started. """ actions_count = 0 try: if self._CheckIfHuntTaskWasAssigned(client_id, rule.hunt_id): logging.info( "Foreman: ignoring hunt %s on client %s: was started " "here before", client_id, rule.hunt_id) else: # hunt_name is only used for legacy hunts. if rule.hunt_name: flow_cls = registry.AFF4FlowRegistry.FlowClassByName( rule.hunt_name) hunt_urn = rdfvalue.RDFURN("aff4:/hunts/%s" % rule.hunt_id) flow_cls.StartClients(hunt_urn, [client_id]) else: try: hunt.StartHuntFlowOnClient(client_id, rule.hunt_id) logging.info("Foreman: Started hunt %s on client %s.", rule.hunt_id, client_id) except flow.CanNotStartFlowWithExistingIdError: logging.info( "Foreman: ignoring hunt %s on client %s: was started " "here before", client_id, rule.hunt_id) actions_count += 1 # There could be all kinds of errors we don't know about when starting the # hunt so we catch everything here. except Exception as e: # pylint: disable=broad-except logging.exception( "Failure running foreman action on client %s: %s", rule.hunt_id, e) return actions_count
def testScheduleHuntRaceCondition(self): client_id = self.SetupClient(0) hunt_id = self._CreateHunt(args=self.GetFileHuntArgs()) original = data_store.REL_DB.delegate.WriteFlowObject def WriteFlowObject(*args, **kwargs): with mock.patch.object(data_store.REL_DB.delegate, "WriteFlowObject", original): try: hunt.StartHuntFlowOnClient(client_id, hunt_id) except Exception as e: raise AssertionError(e) return data_store.REL_DB.WriteFlowObject(*args, **kwargs) # Patch WriteFlowObject to execute another hunt.StartHuntFlowOnClient() for # the same flow and client during the initial StartHuntFlowOnClient(). with mock.patch.object(data_store.REL_DB.delegate, "WriteFlowObject", WriteFlowObject): with self.assertRaises(hunt.flow.CanNotStartFlowWithExistingIdError): hunt.StartHuntFlowOnClient(client_id, hunt_id)
def _RunAction(self, rule, client_id): """Run all the actions specified in the rule. Args: rule: Rule which actions are to be executed. client_id: Id of a client where rule's actions are to be executed. Returns: Number of actions started. """ actions_count = 0 try: if self._CheckIfHuntTaskWasAssigned(client_id, rule.hunt_id): logging.info( "Foreman: ignoring hunt %s on client %s: was started " "here before", rule.hunt_id, client_id) else: try: hunt.StartHuntFlowOnClient(client_id, rule.hunt_id) logging.info("Foreman: Started hunt %s on client %s.", rule.hunt_id, client_id) except flow.CanNotStartFlowWithExistingIdError: logging.info( "Foreman: ignoring hunt %s on client %s: was started " "here before", rule.hunt_id, client_id) actions_count += 1 # There could be all kinds of errors we don't know about when starting the # hunt so we catch everything here. except Exception as e: # pylint: disable=broad-except logging.exception( "Failure running foreman action on client %s: %s", rule.hunt_id, e) return actions_count