def WriteHuntResults(client_id, hunt_id, responses):
"""Writes hunt results from a given client as part of a given hunt."""
if not hunt.IsLegacyHunt(hunt_id):
data_store.REL_DB.WriteFlowResults(responses)
def UpdateFn(hunt_obj):
hunt_obj.num_results += len(responses)
return hunt_obj
hunt_obj = data_store.REL_DB.UpdateHuntObject(hunt_id, UpdateFn)
hunt_obj = hunt.StopHuntIfAverageLimitsExceeded(hunt_obj)
return
hunt_id_urn = rdfvalue.RDFURN("hunts").Add(hunt_id)
msgs = []
for response in responses:
if isinstance(response, rdf_flow_objects.FlowStatus):
continue
msgs.append(
rdf_flows.GrrMessage(payload=response.payload, source=client_id))
with data_store.DB.GetMutationPool() as pool:
for msg in msgs:
hunts_results.HuntResultCollection.StaticAdd(
hunt_id_urn.Add("Results"), msg, mutation_pool=pool)
for msg in msgs:
multi_type_collection.MultiTypeCollection.StaticAdd(
hunt_id_urn.Add("ResultsPerType"), msg, mutation_pool=pool)
stats_collector_instance.Get().IncrementCounter("hunt_results_added",
delta=len(responses))
def ProcessHuntFlowError(flow_obj,
error_message=None,
backtrace=None,
status_msg=None):
"""Processes error and status message for a given hunt-induced flow."""
if not hunt.IsLegacyHunt(flow_obj.parent_hunt_id):
def UpdateFn(hunt_obj):
hunt_obj.num_failed_clients += 1
return hunt_obj
hunt_obj = data_store.REL_DB.UpdateHuntObject(flow_obj.parent_hunt_id,
UpdateFn)
hunt_obj = hunt.StopHuntIfAverageLimitsExceeded(hunt_obj)
return
hunt_urn = rdfvalue.RDFURN("hunts").Add(flow_obj.parent_hunt_id)
client_urn = rdf_client.ClientURN(flow_obj.client_id)
error = rdf_hunts.HuntError(client_id=flow_obj.client_id,
backtrace=backtrace)
if error_message is not None:
error.log_message = error_message
with data_store.DB.GetMutationPool() as pool:
grr_collections.HuntErrorCollection.StaticAdd(
hunt_urn.Add("ErrorClients"), error, mutation_pool=pool)
grr_collections.ClientUrnCollection.StaticAdd(
hunt_urn.Add("CompletedClients"), client_urn, mutation_pool=pool)
if status_msg is not None:
with aff4.FACTORY.OpenWithLock(hunt_urn,
lease_time=_HUNT_LEASE_TIME,
blocking=True) as fd:
fd.GetRunner().SaveResourceUsage(flow_obj.client_id, status_msg)
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 ProcessHuntFlowDone(flow_obj, status_msg=None):
"""Notifis hunt about a given hunt-induced flow completion."""
if not hunt.IsLegacyHunt(flow_obj.parent_hunt_id):
hunt_obj = hunt.StopHuntIfCPUOrNetworkLimitsExceeded(
flow_obj.parent_hunt_id)
hunt.CompleteHuntIfExpirationTimeReached(hunt_obj)
return
hunt_urn = rdfvalue.RDFURN("hunts").Add(flow_obj.parent_hunt_id)
client_urn = rdf_client.ClientURN(flow_obj.client_id)
# Update the counter metrics separately from collections to minimize
# contention.
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)
if flow_obj.num_replies_sent:
fd.context.clients_with_results_count += 1
fd.context.completed_clients_count += 1
fd.context.results_count += flow_obj.num_replies_sent
fd.GetRunner().SaveResourceUsage(flow_obj.client_id, status_msg)
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.RegisterCompletedClient(client_urn)
if flow_obj.num_replies_sent:
fd.RegisterClientWithResults(client_urn)
fd.StopHuntIfAverageLimitsExceeded()
def ProcessHuntFlowError(flow_obj,
error_message=None,
backtrace=None,
status_msg=None):
"""Processes error and status message for a given hunt-induced flow."""
if not hunt.IsLegacyHunt(flow_obj.parent_hunt_id):
hunt.StopHuntIfCPUOrNetworkLimitsExceeded(flow_obj.parent_hunt_id)
return
hunt_urn = rdfvalue.RDFURN("hunts").Add(flow_obj.parent_hunt_id)
client_urn = rdf_client.ClientURN(flow_obj.client_id)
error = rdf_hunts.HuntError(client_id=flow_obj.client_id,
backtrace=backtrace)
if error_message is not None:
error.log_message = error_message
with data_store.DB.GetMutationPool() as pool:
grr_collections.HuntErrorCollection.StaticAdd(
hunt_urn.Add("ErrorClients"), error, mutation_pool=pool)
grr_collections.ClientUrnCollection.StaticAdd(
hunt_urn.Add("CompletedClients"), client_urn, mutation_pool=pool)
if status_msg is not None:
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.GetRunner().SaveResourceUsage(flow_obj.client_id, status_msg)
def WriteAllCrashDetails(client_id,
crash_details,
flow_session_id=None,
hunt_session_id=None,
token=None):
"""Updates the last crash attribute of the client."""
# AFF4.
if data_store.AFF4Enabled():
with aff4.FACTORY.Create(client_id, aff4_grr.VFSGRRClient,
token=token) as client_obj:
client_obj.Set(client_obj.Schema.LAST_CRASH(crash_details))
# Duplicate the crash information in a number of places so we can find it
# easily.
client_urn = rdf_client.ClientURN(client_id)
client_crashes = aff4_grr.VFSGRRClient.CrashCollectionURNForCID(
client_urn)
with data_store.DB.GetMutationPool() as pool:
grr_collections.CrashCollection.StaticAdd(client_crashes,
crash_details,
mutation_pool=pool)
# Relational db.
if data_store.RelationalDBEnabled():
try:
data_store.REL_DB.WriteClientCrashInfo(client_id, crash_details)
except db.UnknownClientError:
pass
if not flow_session_id:
return
if data_store.RelationalDBEnabled():
flow_id = flow_session_id.Basename()
data_store.REL_DB.UpdateFlow(client_id,
flow_id,
client_crash_info=crash_details)
flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
if (flow_obj.parent_hunt_id
and not hunt.IsLegacyHunt(flow_obj.parent_hunt_id)):
hunt.StopHuntIfCrashLimitExceeded(flow_obj.parent_hunt_id)
else:
with aff4.FACTORY.Open(flow_session_id,
flow.GRRFlow,
mode="rw",
age=aff4.NEWEST_TIME,
token=token) as aff4_flow:
aff4_flow.Set(aff4_flow.Schema.CLIENT_CRASH(crash_details))
hunt_session_id = ExtractHuntId(flow_session_id)
if hunt_session_id and hunt_session_id != flow_session_id:
hunt_obj = aff4.FACTORY.Open(hunt_session_id,
aff4_type=implementation.GRRHunt,
mode="rw",
token=token)
hunt_obj.RegisterCrash(crash_details)
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):
hunt.StopHuntIfCrashLimitExceeded(flow_obj.parent_hunt_id)
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 ProcessHuntFlowLog(flow_obj, log_msg):
"""Processes log message from a given hunt-induced flow."""
if not hunt.IsLegacyHunt(flow_obj.parent_hunt_id):
return
hunt_urn = rdfvalue.RDFURN("hunts").Add(flow_obj.parent_hunt_id)
flow_urn = hunt_urn.Add(flow_obj.flow_id)
log_entry = rdf_flows.FlowLog(client_id=flow_obj.client_id,
urn=flow_urn,
flow_name=flow_obj.flow_class_name,
log_message=log_msg)
with data_store.DB.GetMutationPool() as pool:
grr_collections.LogCollection.StaticAdd(hunt_urn.Add("Logs"),
log_entry,
mutation_pool=pool)
def ProcessHuntFlowDone(flow_obj, status_msg=None):
"""Notifis hunt about a given hunt-induced flow completion."""
if not hunt.IsLegacyHunt(flow_obj.parent_hunt_id):
resources = rdf_client_stats.ClientResources(
client_id=flow_obj.client_id,
session_id=flow_obj.flow_id,
cpu_usage=rdf_client_stats.CpuSeconds(
user_cpu_time=status_msg.cpu_time_used.user_cpu_time,
system_cpu_time=status_msg.cpu_time_used.system_cpu_time),
network_bytes_sent=status_msg.network_bytes_sent)
def UpdateFn(hunt_obj):
hunt_obj.num_successful_clients += 1
if flow_obj.num_replies_sent:
hunt_obj.num_clients_with_results += 1
hunt_obj.client_resources_stats.RegisterResources(resources)
return hunt_obj
hunt_obj = data_store.REL_DB.UpdateHuntObject(flow_obj.parent_hunt_id,
UpdateFn)
hunt_obj = hunt.StopHuntIfAverageLimitsExceeded(hunt_obj)
hunt.CompleteHuntIfExpirationTimeReached(hunt_obj)
return
hunt_urn = rdfvalue.RDFURN("hunts").Add(flow_obj.parent_hunt_id)
client_urn = rdf_client.ClientURN(flow_obj.client_id)
with aff4.FACTORY.OpenWithLock(hunt_urn,
lease_time=_HUNT_LEASE_TIME,
blocking=True) as fd:
# Legacy AFF4 code expects token to be set.
fd.token = access_control.ACLToken(username=fd.creator)
fd.RegisterCompletedClient(client_urn)
if flow_obj.num_replies_sent:
fd.RegisterClientWithResults(client_urn)
fd.context.clients_with_results_count += 1
fd.context.completed_clients_count += 1
fd.context.results_count += flow_obj.num_replies_sent
fd.GetRunner().SaveResourceUsage(flow_obj.client_id, status_msg)
fd.StopHuntIfAverageLimitsExceeded()
def _CheckIfHuntTaskWasAssigned(self, client_id, hunt_id):
"""Will return True if hunt's task was assigned to this client before."""
if data_store.RelationalDBEnabled():
flow_id = hunt_id
if hunt.IsLegacyHunt(hunt_id):
# Strip "H:" prefix.
flow_id = flow_id[2:]
try:
data_store.REL_DB.ReadFlowObject(client_id, flow_id)
return True
except db.UnknownFlowError:
pass
else:
client_urn = rdfvalue.RDFURN(client_id)
for _ in aff4.FACTORY.Stat([
client_urn.Add("flows/%s:hunt" % rdfvalue.RDFURN(hunt_id).Basename())
]):
return True
return False