def SaveResourceUsage(self, status):
"""Method to tally resources."""
user_cpu = status.cpu_time_used.user_cpu_time
system_cpu = status.cpu_time_used.system_cpu_time
self.rdf_flow.cpu_time_used.user_cpu_time += user_cpu
self.rdf_flow.cpu_time_used.system_cpu_time += system_cpu
self.rdf_flow.network_bytes_sent += status.network_bytes_sent
if self.rdf_flow.cpu_limit:
user_cpu_total = self.rdf_flow.cpu_time_used.user_cpu_time
system_cpu_total = self.rdf_flow.cpu_time_used.system_cpu_time
if self.rdf_flow.cpu_limit < (user_cpu_total + system_cpu_total):
# We have exceeded our limit, stop this flow.
raise flow.FlowResourcesExceededError(
"CPU limit exceeded for {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
if (self.rdf_flow.network_bytes_limit
and self.rdf_flow.network_bytes_limit <
self.rdf_flow.network_bytes_sent):
# We have exceeded our byte limit, stop this flow.
raise flow.FlowResourcesExceededError(
"Network bytes limit exceeded {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
def SaveResourceUsage(self, status: rdf_flows.GrrStatus) -> None:
"""Method to tally resources."""
user_cpu = status.cpu_time_used.user_cpu_time
system_cpu = status.cpu_time_used.system_cpu_time
self.rdf_flow.cpu_time_used.user_cpu_time += user_cpu
self.rdf_flow.cpu_time_used.system_cpu_time += system_cpu
self.rdf_flow.network_bytes_sent += status.network_bytes_sent
if not self.rdf_flow.runtime_us:
self.rdf_flow.runtime_us = rdfvalue.Duration(0)
if status.runtime_us:
self.rdf_flow.runtime_us += status.runtime_us
if self.rdf_flow.cpu_limit:
user_cpu_total = self.rdf_flow.cpu_time_used.user_cpu_time
system_cpu_total = self.rdf_flow.cpu_time_used.system_cpu_time
if self.rdf_flow.cpu_limit < (user_cpu_total + system_cpu_total):
# We have exceeded our CPU time limit, stop this flow.
raise flow.FlowResourcesExceededError(
"CPU limit exceeded for {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
if (self.rdf_flow.network_bytes_limit
and self.rdf_flow.network_bytes_limit <
self.rdf_flow.network_bytes_sent):
# We have exceeded our byte limit, stop this flow.
raise flow.FlowResourcesExceededError(
"Network bytes limit exceeded {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
if (self.rdf_flow.runtime_limit_us
and self.rdf_flow.runtime_limit_us < self.rdf_flow.runtime_us):
raise flow.FlowResourcesExceededError(
"Runtime limit exceeded {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
def CallClient(self,
action_cls: Type[server_stubs.ClientActionStub],
request: Optional[rdfvalue.RDFValue] = None,
next_state: Optional[str] = None,
callback_state: Optional[str] = None,
request_data: Optional[Mapping[str, Any]] = None,
**kwargs: Any):
"""Calls the client asynchronously.
This sends a message to the client to invoke an Action. The run action may
send back many responses that will be queued by the framework until a status
message is sent by the client. The status message will cause the entire
transaction to be committed to the specified state.
Args:
action_cls: The function to call on the client.
request: The request to send to the client. If not specified, we create a
new RDFValue using the kwargs.
next_state: The state in this flow, that responses to this message should
go to.
callback_state: (optional) The state to call whenever a new response is
arriving.
request_data: A dict which will be available in the RequestState
protobuf. The Responses object maintains a reference to this protobuf
for use in the execution of the state method. (so you can access this
data by responses.request).
**kwargs: These args will be used to construct the client action argument
rdfvalue.
Raises:
ValueError: The request passed to the client does not have the correct
type.
"""
try:
action_identifier = action_registry.ID_BY_ACTION_STUB[action_cls]
except KeyError:
raise ValueError("Action class %s not known." % action_cls)
if action_cls.in_rdfvalue is None:
if request:
raise ValueError("Client action %s does not expect args." %
action_cls)
else:
if request is None:
# Create a new rdf request.
request = action_cls.in_rdfvalue(**kwargs)
else:
# Verify that the request type matches the client action requirements.
if not isinstance(request, action_cls.in_rdfvalue):
raise ValueError("Client action expected %s but got %s" %
(action_cls.in_rdfvalue, type(request)))
outbound_id = self.GetNextOutboundId()
# Create a flow request.
flow_request = rdf_flow_objects.FlowRequest(
client_id=self.rdf_flow.client_id,
flow_id=self.rdf_flow.flow_id,
request_id=outbound_id,
next_state=next_state,
callback_state=callback_state)
if request_data is not None:
flow_request.request_data = rdf_protodict.Dict().FromDict(
request_data)
cpu_limit_ms = None
network_bytes_limit = None
if self.rdf_flow.cpu_limit:
cpu_usage = self.rdf_flow.cpu_time_used
cpu_limit_ms = 1000 * max(
self.rdf_flow.cpu_limit - cpu_usage.user_cpu_time -
cpu_usage.system_cpu_time, 0)
if cpu_limit_ms == 0:
raise flow.FlowResourcesExceededError(
"CPU limit exceeded for {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
if self.rdf_flow.network_bytes_limit:
network_bytes_limit = max(
self.rdf_flow.network_bytes_limit -
self.rdf_flow.network_bytes_sent, 0)
if network_bytes_limit == 0:
raise flow.FlowResourcesExceededError(
"Network limit exceeded for {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
runtime_limit_us = self.rdf_flow.runtime_limit_us
if runtime_limit_us and self.rdf_flow.runtime_us:
if self.rdf_flow.runtime_us < runtime_limit_us:
runtime_limit_us -= self.rdf_flow.runtime_us
else:
raise flow.FlowResourcesExceededError(
"Runtime limit exceeded for {} {}.".format(
self.rdf_flow.flow_class_name, self.rdf_flow.flow_id))
client_action_request = rdf_flows.ClientActionRequest(
client_id=self.rdf_flow.client_id,
flow_id=self.rdf_flow.flow_id,
request_id=outbound_id,
action_identifier=action_identifier,
action_args=request,
cpu_limit_ms=cpu_limit_ms,
network_bytes_limit=network_bytes_limit,
runtime_limit_us=runtime_limit_us)
self.flow_requests.append(flow_request)
self.client_action_requests.append(client_action_request)
