def testPendingFlowTermination(self): client_mock = ClientMock() flow_id = flow.StartFlow(flow_cls=ParentFlow, client_id=self.client_id) flow_obj = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id) self.assertEqual(flow_obj.flow_state, "RUNNING") pending_termination = rdf_flow_objects.PendingFlowTermination( reason="testing") data_store.REL_DB.UpdateFlow(self.client_id, flow_id, pending_termination=pending_termination) with flow_test_lib.TestWorker() as worker: with test_lib.SuppressLogs(): flow_test_lib.RunFlow(self.client_id, flow_id, client_mock=client_mock, worker=worker, check_flow_errors=False) flow_obj = data_store.REL_DB.ReadFlowObject( self.client_id, flow_id) self.assertEqual(flow_obj.flow_state, "ERROR") self.assertEqual(flow_obj.error_message, "testing")
def StopHunt(hunt_id, reason=None): """Stops a hunt with a given id.""" def UpdateFn(h): if h.hunt_state not in [h.HuntState.STARTED, h.HuntState.PAUSED]: raise OnlyStartedOrPausedHuntCanBeStoppedError(h) h.hunt_state = h.HuntState.STOPPED if reason is not None: h.hunt_state_comment = reason return h # If the hunt was not started or paused, the exception from UpdateFn is # guaranteed to be propagated by UpdateHuntObject implementation. hunt_obj = data_store.REL_DB.UpdateHuntObject(hunt_id, UpdateFn) data_store.REL_DB.RemoveForemanRule(hunt_id=hunt_obj.hunt_id) flows = data_store.REL_DB.ReadHuntFlows(hunt_obj.hunt_id, 0, sys.maxsize) data_store.REL_DB.UpdateFlows( [(f.client_id, f.flow_id) for f in flows], pending_termination=rdf_flow_objects.PendingFlowTermination( reason="Parent hunt stopped.")) if (reason is not None and hunt_obj.creator not in aff4_users.GRRUser.SYSTEM_USERS): notification.Notify( hunt_obj.creator, rdf_objects.UserNotification.Type.TYPE_HUNT_STOPPED, reason, rdf_objects.ObjectReference( reference_type=rdf_objects.ObjectReference.Type.HUNT, hunt=rdf_objects.HuntReference(hunt_id=hunt_obj.hunt_id))) return hunt_obj
def testPendingTerminationUpdate(self): client_id, flow_id = self._SetupClientAndFlow() pending_termination = rdf_flow_objects.PendingFlowTermination( reason="test") self.db.UpdateFlow(client_id, flow_id, pending_termination=pending_termination) read_flow = self.db.ReadFlowObject(client_id, flow_id) self.assertEqual(read_flow.pending_termination, pending_termination)
def MarkForTermination(cls, flow_urn, reason=None, mutation_pool=None): """Mark the flow for termination as soon as any of its states are called.""" # Doing a blind write here using low-level data store API. Accessing # the flow via AFF4 is not really possible here, because it forces a state # to be written in Close() method. if mutation_pool is None: raise ValueError("Mutation pool can't be none.") mutation_pool.Set( flow_urn, cls.SchemaCls.PENDING_TERMINATION.predicate, rdf_flow_objects.PendingFlowTermination(reason=reason), replace=False)
def _StopRelational(self, reason=None): super(GenericHunt, self).Stop(reason=reason) started_flows = grr_collections.RDFUrnCollection( self.started_flows_collection_urn) client_id_flow_id_pairs = [] for flow_urn in started_flows: components = flow_urn.Split() client_id_flow_id_pairs.append((components[0], components[2])) data_store.REL_DB.UpdateFlows( client_id_flow_id_pairs, pending_termination=rdf_flow_objects.PendingFlowTermination( reason="Parent hunt stopped."))