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."))
