def testDrainTaskSchedulerQueue(self):
client_id = u"C.1234567890123456"
flow_id = flow.RandomFlowId()
data_store.REL_DB.WriteClientMetadata(client_id,
fleetspeak_enabled=False)
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow_id,
create_time=rdfvalue.RDFDatetime.Now())
data_store.REL_DB.WriteFlowObject(rdf_flow)
action_requests = []
for i in range(3):
data_store.REL_DB.WriteFlowRequests([
rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=i)
])
action_requests.append(
rdf_flows.ClientActionRequest(client_id=client_id,
flow_id=flow_id,
request_id=i,
action_identifier="WmiQuery"))
data_store.REL_DB.WriteClientActionRequests(action_requests)
server = TestServer()
res = server.DrainTaskSchedulerQueueForClient(
rdfvalue.RDFURN(client_id))
msgs = [
rdf_flow_objects.GRRMessageFromClientActionRequest(r)
for r in action_requests
]
for r in res:
r.task_id = 0
for m in msgs:
m.task_id = 0
self.assertItemsEqual(res, msgs)
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)
