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)