def LeaseFlowForProcessing(self, client_id, flow_id, processing_time):
"""Marks a flow as being processed on this worker and returns it."""
rdf_flow = self.ReadFlowObject(client_id, flow_id)
if rdf_flow.parent_hunt_id:
rdf_hunt = self.ReadHuntObject(rdf_flow.parent_hunt_id)
if not rdf_hunt_objects.IsHuntSuitableForFlowProcessing(
rdf_hunt.hunt_state):
raise db.ParentHuntIsNotRunningError(client_id, flow_id,
rdf_hunt.hunt_id,
rdf_hunt.hunt_state)
now = rdfvalue.RDFDatetime.Now()
if rdf_flow.processing_on and rdf_flow.processing_deadline > now:
raise ValueError("Flow %s on client %s is already being processed." %
(flow_id, client_id))
processing_deadline = now + processing_time
process_id_string = utils.ProcessIdString()
self.UpdateFlow(
client_id,
flow_id,
processing_on=process_id_string,
processing_since=now,
processing_deadline=processing_deadline)
rdf_flow.processing_on = process_id_string
rdf_flow.processing_since = now
rdf_flow.processing_deadline = processing_deadline
return rdf_flow
def ReadFlowForProcessing(self, client_id, flow_id, processing_time):
"""Marks a flow as being processed on this worker and returns it."""
rdf_flow = self.ReadFlowObject(client_id, flow_id)
# TODO(user): remove the check for a legacy hunt prefix as soon as
# AFF4 is gone.
if rdf_flow.parent_hunt_id and not rdf_flow.parent_hunt_id.startswith("H:"):
rdf_hunt = self.ReadHuntObject(rdf_flow.parent_hunt_id)
if not rdf_hunt_objects.IsHuntSuitableForFlowProcessing(
rdf_hunt.hunt_state):
raise db.ParentHuntIsNotRunningError(
client_id, flow_id, rdf_hunt.hunt_id, rdf_hunt.hunt_state)
now = rdfvalue.RDFDatetime.Now()
if rdf_flow.processing_on and rdf_flow.processing_deadline > now:
raise ValueError("Flow %s on client %s is already being processed." %
(client_id, flow_id))
processing_deadline = now + processing_time
process_id_string = utils.ProcessIdString()
self.UpdateFlow(
client_id,
flow_id,
processing_on=process_id_string,
processing_since=now,
processing_deadline=processing_deadline)
rdf_flow.processing_on = process_id_string
rdf_flow.processing_since = now
rdf_flow.processing_deadline = processing_deadline
return rdf_flow
def StartHuntFlowOnClient(client_id, hunt_id):
"""Starts a flow corresponding to a given hunt on a given client."""
hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id)
hunt_obj = CompleteHuntIfExpirationTimeReached(hunt_obj)
# There may be a little race between foreman rules being removed and
# foreman scheduling a client on an (already) paused hunt. Making sure
# we don't lose clients in such a race by accepting clients for paused
# hunts.
if not rdf_hunt_objects.IsHuntSuitableForFlowProcessing(
hunt_obj.hunt_state):
return
if hunt_obj.args.hunt_type == hunt_obj.args.HuntType.STANDARD:
hunt_args = hunt_obj.args.standard
if hunt_obj.client_rate > 0:
# Given that we use caching in _GetNumClients and hunt_obj may be updated
# in another process, we have to account for cases where num_clients_diff
# may go below 0.
num_clients_diff = max(
0,
_GetNumClients(hunt_obj.hunt_id) -
hunt_obj.num_clients_at_start_time)
next_client_due_msecs = int(num_clients_diff /
hunt_obj.client_rate * 60e6)
start_at = rdfvalue.RDFDatetime.FromMicrosecondsSinceEpoch(
hunt_obj.last_start_time.AsMicrosecondsSinceEpoch() +
next_client_due_msecs)
else:
start_at = None
# TODO(user): remove client_rate support when AFF4 is gone.
# In REL_DB always work as if client rate is 0.
flow_cls = registry.FlowRegistry.FlowClassByName(hunt_args.flow_name)
flow_args = hunt_args.flow_args if hunt_args.HasField(
"flow_args") else None
flow.StartFlow(
client_id=client_id,
creator=hunt_obj.creator,
cpu_limit=hunt_obj.per_client_cpu_limit,
network_bytes_limit=hunt_obj.per_client_network_bytes_limit,
flow_cls=flow_cls,
flow_args=flow_args,
start_at=start_at,
parent_hunt_id=hunt_id)
if hunt_obj.client_limit:
if _GetNumClients(hunt_obj.hunt_id) >= hunt_obj.client_limit:
PauseHunt(hunt_id)
elif hunt_obj.args.hunt_type == hunt_obj.args.HuntType.VARIABLE:
raise NotImplementedError()
else:
raise UnknownHuntTypeError("Can't determine hunt type when starting "
"hunt %s on client %s." %
(client_id, hunt_id))
