def _SendTerminationMessage(self, status=None):
"""This notifies the parent flow of our termination."""
if not self.runner_args.request_state.session_id:
# No parent flow, nothing to do here.
return
if status is None:
status = rdf_flows.GrrStatus()
client_resources = self.context.client_resources
user_cpu = client_resources.cpu_usage.user_cpu_time
sys_cpu = client_resources.cpu_usage.system_cpu_time
status.cpu_time_used.user_cpu_time = user_cpu
status.cpu_time_used.system_cpu_time = sys_cpu
status.network_bytes_sent = self.context.network_bytes_sent
status.child_session_id = self.session_id
request_state = self.runner_args.request_state
request_state.response_count += 1
# Make a response message
msg = rdf_flows.GrrMessage(
session_id=request_state.session_id,
request_id=request_state.id,
response_id=request_state.response_count,
auth_state=rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED,
type=rdf_flows.GrrMessage.Type.STATUS,
payload=status)
# Queue the response now
self.queue_manager.QueueResponse(msg)
self.QueueNotification(session_id=request_state.session_id)
def testUnauthenticated(self):
"""What happens if an unauthenticated message is sent to the client?
RuntimeError needs to be issued, and the client needs to send a
GrrStatus message with the traceback in it.
"""
# Push a request on it
message = rdf_flows.GrrMessage(
name="MockAction",
session_id=self.session_id,
auth_state=rdf_flows.GrrMessage.AuthorizationState.UNAUTHENTICATED,
request_id=1,
generate_task_id=True)
self.context.HandleMessage(message)
# We expect to receive an GrrStatus to indicate an exception was
# raised:
# Check the response - one data and one status
message_list = self.context.Drain().job
self.assertEqual(len(message_list), 1)
self.assertEqual(message_list[0].session_id, self.session_id)
self.assertEqual(message_list[0].response_id, 1)
status = rdf_flows.GrrStatus(message_list[0].payload)
self.assertIn("not Authenticated", status.error_message)
self.assertIn("RuntimeError", status.error_message)
self.assertNotEqual(status.status,
rdf_flows.GrrStatus.ReturnedStatus.OK)
def RunAction(self, action_cls, arg=None, grr_worker=None):
if arg is None:
arg = rdf_flows.GrrMessage()
self.results = []
action = self._GetActionInstance(action_cls, grr_worker=grr_worker)
action.status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
action.Run(arg)
return self.results
def testEqualTimestampNotifications(self):
frontend_server = frontend_lib.FrontEndServer(
certificate=config.CONFIG["Frontend.certificate"],
private_key=config.CONFIG["PrivateKeys.server_key"],
message_expiry_time=100,
threadpool_prefix="notification-test")
# This schedules 10 requests.
session_id = flow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
# We pretend that the client processed all the 10 requests at once and
# sends the replies in a single http poll.
messages = [
rdf_flows.GrrMessage(
request_id=i,
response_id=1,
session_id=session_id,
payload=rdf_protodict.DataBlob(string="test%s" % i),
auth_state="AUTHENTICATED",
generate_task_id=True) for i in range(1, 11)
]
status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
statuses = [
rdf_flows.GrrMessage(request_id=i,
response_id=2,
session_id=session_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS,
auth_state="AUTHENTICATED",
generate_task_id=True) for i in range(1, 11)
]
frontend_server.ReceiveMessages(self.client_id, messages + statuses)
with queue_manager.QueueManager(token=self.token) as q:
all_notifications = q.GetNotificationsByPriorityForAllShards(
rdfvalue.RDFURN("aff4:/F"))
medium_priority = rdf_flows.GrrNotification.Priority.MEDIUM_PRIORITY
medium_notifications = all_notifications[medium_priority]
my_notifications = [
n for n in medium_notifications if n.session_id == session_id
]
# There must not be more than one notification.
self.assertEqual(len(my_notifications), 1)
notification = my_notifications[0]
self.assertEqual(notification.first_queued, notification.timestamp)
self.assertEqual(notification.last_status, 10)
def SendOKStatus(self, response_id, session_id):
"""Send a message to the flow."""
message = rdf_flows.GrrMessage(
request_id=1,
response_id=response_id,
session_id=session_id,
type=rdf_flows.GrrMessage.Type.STATUS,
auth_state=rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED)
status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
message.payload = status
self.SendMessage(message)
def Run(self, unused_arg):
"""Run the kill."""
# Send a message back to the service to say that we are about to shutdown.
reply = rdf_flows.GrrStatus(status=rdf_flows.GrrStatus.ReturnedStatus.OK)
# Queue up the response message, jump the queue.
self.SendReply(
reply,
message_type=rdf_flows.GrrMessage.Type.STATUS,
priority=rdf_flows.GrrMessage.Priority.HIGH_PRIORITY + 1)
# Give the http thread some time to send the reply.
self.grr_worker.Sleep(10)
# Die ourselves.
logging.info("Dying on request.")
os._exit(242) # pylint: disable=protected-access
def Error(self, backtrace, client_id=None, status_code=None):
"""Terminates this flow with an error."""
try:
self.queue_manager.DestroyFlowStates(self.session_id)
except queue_manager.MoreDataException:
pass
if not self.IsRunning():
return
# Set an error status
reply = rdf_flows.GrrStatus()
if status_code is None:
reply.status = rdf_flows.GrrStatus.ReturnedStatus.GENERIC_ERROR
else:
reply.status = status_code
client_id = client_id or self.runner_args.client_id
if backtrace:
reply.error_message = backtrace
logging.error("Error in flow %s (%s). Trace: %s", self.session_id,
client_id, backtrace)
self.context.backtrace = backtrace
else:
logging.error("Error in flow %s (%s).", self.session_id, client_id)
self._SendTerminationMessage(reply)
self.context.state = rdf_flow_runner.FlowContext.State.ERROR
if self.ShouldSendNotifications():
flow_ref = None
if client_id:
flow_ref = rdf_objects.FlowReference(
client_id=client_id.Basename(),
flow_id=self.session_id.Basename())
notification_lib.Notify(
self.token.username,
rdf_objects.UserNotification.Type.TYPE_FLOW_RUN_FAILED,
"Flow (%s) terminated due to error" % self.session_id,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.FLOW,
flow=flow_ref))
self.flow_obj.Flush()
def __init__(self, grr_worker=None):
"""Initializes the action plugin.
Args:
grr_worker: The grr client worker object which may be used to
e.g. send new actions on.
"""
self.grr_worker = grr_worker
self.response_id = INITIAL_RESPONSE_ID
self.cpu_used = None
self.nanny_controller = None
self.status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
self._last_gc_run = rdfvalue.RDFDatetime.Now()
self._gc_frequency = config.CONFIG["Client.gc_frequency"]
self.proc = psutil.Process()
self.cpu_start = self.proc.cpu_times()
self.cpu_limit = rdf_flows.GrrMessage().cpu_limit
def testHandleError(self):
"""Test handling of a request which raises."""
# Push a request on it
message = rdf_flows.GrrMessage(
name="RaiseAction",
session_id=self.session_id,
auth_state=rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED,
request_id=1,
generate_task_id=True)
self.context.HandleMessage(message)
# Check the response - one data and one status
message_list = self.context.Drain().job
self.assertEqual(message_list[0].session_id, self.session_id)
self.assertEqual(message_list[0].response_id, 1)
status = rdf_flows.GrrStatus(message_list[0].payload)
self.assertIn("RuntimeError", status.error_message)
self.assertNotEqual(status.status,
rdf_flows.GrrStatus.ReturnedStatus.OK)
def StatFile(self, args):
"""StatFile action mock."""
response = rdf_client.StatEntry(
pathspec=args.pathspec,
st_mode=33184,
st_ino=1063090,
st_dev=64512L,
st_nlink=1,
st_uid=139592,
st_gid=5000,
st_size=len(self.data),
st_atime=1336469177,
st_mtime=1336129892,
st_ctime=1336129892)
self.responses += 1
self.count += 1
# Create status message to report sample resource usage
status = rdf_flows.GrrStatus(status=rdf_flows.GrrStatus.ReturnedStatus.OK)
if self.user_cpu_time is None:
status.cpu_time_used.user_cpu_time = self.responses
else:
status.cpu_time_used.user_cpu_time = self.user_cpu_time
if self.system_cpu_time is None:
status.cpu_time_used.system_cpu_time = self.responses * 2
else:
status.cpu_time_used.system_cpu_time = self.system_cpu_time
if self.network_bytes_sent is None:
status.network_bytes_sent = self.responses * 3
else:
status.network_bytes_sent = self.network_bytes_sent
# Every "failrate" client does not have this file.
if self.count == self.failrate:
self.count = 0
return [status]
return [response, status]
def SendResponse(self,
session_id,
data,
client_id=None,
well_known=False,
request_id=None):
if not isinstance(data, rdfvalue.RDFValue):
data = rdf_protodict.DataBlob(string=data)
if well_known:
request_id, response_id = 0, 12345
else:
request_id, response_id = request_id or 1, 1
with queue_manager.QueueManager(token=self.token) as flow_manager:
flow_manager.QueueResponse(
rdf_flows.GrrMessage(source=client_id,
session_id=session_id,
payload=data,
request_id=request_id,
auth_state="AUTHENTICATED",
response_id=response_id))
if not well_known:
# For normal flows we have to send a status as well.
flow_manager.QueueResponse(
rdf_flows.GrrMessage(
source=client_id,
session_id=session_id,
payload=rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK),
request_id=request_id,
response_id=response_id + 1,
auth_state="AUTHENTICATED",
type=rdf_flows.GrrMessage.Type.STATUS))
flow_manager.QueueNotification(session_id=session_id,
last_status=request_id)
timestamp = flow_manager.frozen_timestamp
return timestamp
def ReceiveMessages(self, client_id, messages):
"""Receives and processes the messages from the source.
For each message we update the request object, and place the
response in that request's queue. If the request is complete, we
send a message to the worker.
Args:
client_id: The client which sent the messages.
messages: A list of GrrMessage RDFValues.
"""
now = time.time()
with queue_manager.QueueManager(token=self.token) as manager:
for session_id, msgs in utils.GroupBy(
messages, operator.attrgetter("session_id")).iteritems():
# Remove and handle messages to WellKnownFlows
leftover_msgs = self.HandleWellKnownFlows(msgs)
unprocessed_msgs = []
for msg in leftover_msgs:
if (msg.auth_state == msg.AuthorizationState.AUTHENTICATED
or msg.session_id
== self.unauth_allowed_session_id):
unprocessed_msgs.append(msg)
if len(unprocessed_msgs) < len(leftover_msgs):
logging.info("Dropped %d unauthenticated messages for %s",
len(leftover_msgs) - len(unprocessed_msgs),
client_id)
if not unprocessed_msgs:
continue
for msg in unprocessed_msgs:
manager.QueueResponse(msg)
for msg in unprocessed_msgs:
# Messages for well known flows should notify even though they don't
# have a status.
if msg.request_id == 0:
manager.QueueNotification(session_id=msg.session_id,
priority=msg.priority)
# Those messages are all the same, one notification is enough.
break
elif msg.type == rdf_flows.GrrMessage.Type.STATUS:
# If we receive a status message from the client it means the client
# has finished processing this request. We therefore can de-queue it
# from the client queue. msg.task_id will raise if the task id is
# not set (message originated at the client, there was no request on
# the server), so we have to check .HasTaskID() first.
if msg.HasTaskID():
manager.DeQueueClientRequest(msg)
manager.QueueNotification(session_id=msg.session_id,
priority=msg.priority,
last_status=msg.request_id)
stat = rdf_flows.GrrStatus(msg.payload)
if stat.status == rdf_flows.GrrStatus.ReturnedStatus.CLIENT_KILLED:
# A client crashed while performing an action, fire an event.
crash_details = rdf_client.ClientCrash(
client_id=client_id,
session_id=session_id,
backtrace=stat.backtrace,
crash_message=stat.error_message,
nanny_status=stat.nanny_status,
timestamp=rdfvalue.RDFDatetime.Now())
events.Events.PublishEvent("ClientCrash",
crash_details,
token=self.token)
logging.debug("Received %s messages from %s in %s sec", len(messages),
client_id,
time.time() - now)
def testEnums(self):
"""Check that enums are wrapped in a descriptor class."""
sample = rdf_flows.GrrStatus()
self.assertEqual(str(sample.status), "OK")
def CallState(self,
messages=None,
next_state="",
request_data=None,
start_time=None):
"""This method is used to schedule a new state on a different worker.
This is basically the same as CallFlow() except we are calling
ourselves. The state will be invoked in a later time and receive all the
messages we send.
Args:
messages: A list of rdfvalues to send. If the last one is not a
GrrStatus, we append an OK Status.
next_state: The state in this flow to be invoked with the responses.
request_data: Any dict provided here 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).
start_time: Start the flow at this time. This Delays notification for
flow processing into the future. Note that the flow may still be
processed earlier if there are client responses waiting.
Raises:
FlowRunnerError: if the next state is not valid.
"""
if messages is None:
messages = []
# Check if the state is valid
if not getattr(self.flow_obj, next_state):
raise FlowRunnerError("Next state %s is invalid.")
# Queue the response message to the parent flow
request_state = rdf_flow_runner.RequestState(
id=self.GetNextOutboundId(),
session_id=self.context.session_id,
client_id=self.runner_args.client_id,
next_state=next_state)
if request_data:
request_state.data = rdf_protodict.Dict().FromDict(request_data)
self.QueueRequest(request_state, timestamp=start_time)
# Add the status message if needed.
if not messages or not isinstance(messages[-1], rdf_flows.GrrStatus):
messages.append(rdf_flows.GrrStatus())
# Send all the messages
for i, payload in enumerate(messages):
if isinstance(payload, rdfvalue.RDFValue):
msg = rdf_flows.GrrMessage(
session_id=self.session_id,
request_id=request_state.id,
response_id=1 + i,
auth_state=rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED,
payload=payload,
type=rdf_flows.GrrMessage.Type.MESSAGE)
if isinstance(payload, rdf_flows.GrrStatus):
msg.type = rdf_flows.GrrMessage.Type.STATUS
else:
raise FlowRunnerError(
"Bad message %s of type %s." % (payload, type(payload)))
self.QueueResponse(msg, start_time)
# Notify the worker about it.
self.QueueNotification(session_id=self.session_id, timestamp=start_time)
def testNoValidStatusRaceIsResolved(self):
# This tests for the regression of a long standing race condition we saw
# where notifications would trigger the reading of another request that
# arrives later but wasn't completely written to the database yet.
# Timestamp based notification handling should eliminate this bug.
# We need a random flow object for this test.
session_id = flow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
worker_obj = worker_lib.GRRWorker(token=self.token)
manager = queue_manager.QueueManager(token=self.token)
manager.DeleteNotification(session_id)
manager.Flush()
# We have a first request that is complete (request_id 1, response_id 1).
self.SendResponse(session_id, "Response 1")
# However, we also have request #2 already coming in. The race is that
# the queue manager might write the status notification to
# session_id/state as "status:00000002" but not the status response
# itself yet under session_id/state/request:00000002
request_id = 2
response_id = 1
flow_manager = queue_manager.QueueManager(token=self.token)
flow_manager.FreezeTimestamp()
flow_manager.QueueResponse(
rdf_flows.GrrMessage(
source=self.client_id,
session_id=session_id,
payload=rdf_protodict.DataBlob(string="Response 2"),
request_id=request_id,
auth_state="AUTHENTICATED",
response_id=response_id))
status = rdf_flows.GrrMessage(
source=self.client_id,
session_id=session_id,
payload=rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK),
request_id=request_id,
response_id=response_id + 1,
auth_state="AUTHENTICATED",
type=rdf_flows.GrrMessage.Type.STATUS)
# Now we write half the status information.
data_store.DB.StoreRequestsAndResponses(new_responses=[(status, None)])
# We make the race even a bit harder by saying the new notification gets
# written right before the old one gets deleted. If we are not careful here,
# we delete the new notification as well and the flow becomes stuck.
# pylint: disable=invalid-name
def WriteNotification(self, arg_session_id, start=None, end=None):
if arg_session_id == session_id:
flow_manager.QueueNotification(session_id=arg_session_id)
flow_manager.Flush()
self.DeleteNotification.old_target(self,
arg_session_id,
start=start,
end=end)
# pylint: enable=invalid-name
with utils.Stubber(queue_manager.QueueManager, "DeleteNotification",
WriteNotification):
# This should process request 1 but not touch request 2.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
flow_obj = aff4.FACTORY.Open(session_id, token=self.token)
self.assertFalse(flow_obj.context.backtrace)
self.assertNotEqual(flow_obj.context.state,
rdf_flow_runner.FlowContext.State.ERROR)
request_data = data_store.DB.ReadResponsesForRequestId(session_id, 2)
request_data.sort(key=lambda msg: msg.response_id)
self.assertEqual(len(request_data), 2)
# Make sure the status and the original request are still there.
self.assertEqual(request_data[0].args_rdf_name, "DataBlob")
self.assertEqual(request_data[1].args_rdf_name, "GrrStatus")
# But there is nothing for request 1.
request_data = data_store.DB.ReadResponsesForRequestId(session_id, 1)
self.assertEqual(request_data, [])
# The notification for request 2 should have survived.
with queue_manager.QueueManager(token=self.token) as manager:
notifications = manager.GetNotifications(queues.FLOWS)
self.assertEqual(len(notifications), 1)
notification = notifications[0]
self.assertEqual(notification.session_id, session_id)
self.assertEqual(notification.timestamp,
flow_manager.frozen_timestamp)
self.assertEqual(RESULTS, ["Response 1"])
# The last missing piece of request 2 is the actual status message.
flow_manager.QueueResponse(status)
flow_manager.Flush()
# Now make sure request 2 runs as expected.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
self.assertEqual(RESULTS, ["Response 1", "Response 2"])
def Next(self):
"""Grab tasks for us from the server's queue."""
with queue_manager.QueueManager(token=self.token) as manager:
request_tasks = manager.QueryAndOwn(self.client_id.Queue(),
limit=1,
lease_seconds=10000)
request_tasks.extend(self._mock_task_queue)
self._mock_task_queue[:] = [] # Clear the referenced list.
for message in request_tasks:
status = None
response_id = 1
# Collect all responses for this message from the client mock
try:
if hasattr(self.client_mock, "HandleMessage"):
responses = self.client_mock.HandleMessage(message)
else:
self.client_mock.message = message
responses = getattr(self.client_mock,
message.name)(message.payload)
if not responses:
responses = []
logging.info(
"Called client action %s generating %s responses",
message.name,
len(responses) + 1)
if self.status_message_enforced:
status = rdf_flows.GrrStatus()
except Exception as e: # pylint: disable=broad-except
logging.exception("Error %s occurred in client", e)
# Error occurred.
responses = []
if self.status_message_enforced:
error_message = str(e)
status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.
GENERIC_ERROR)
# Invalid action mock is usually expected.
if error_message != "Invalid Action Mock.":
status.backtrace = traceback.format_exc()
status.error_message = error_message
# Now insert those on the flow state queue
for response in responses:
if isinstance(response, rdf_flows.GrrStatus):
msg_type = rdf_flows.GrrMessage.Type.STATUS
self.AddResourceUsage(response)
response = rdf_flows.GrrMessage(
session_id=message.session_id,
name=message.name,
response_id=response_id,
request_id=message.request_id,
payload=response,
type=msg_type)
elif isinstance(response, rdf_client.Iterator):
msg_type = rdf_flows.GrrMessage.Type.ITERATOR
response = rdf_flows.GrrMessage(
session_id=message.session_id,
name=message.name,
response_id=response_id,
request_id=message.request_id,
payload=response,
type=msg_type)
elif not isinstance(response, rdf_flows.GrrMessage):
msg_type = rdf_flows.GrrMessage.Type.MESSAGE
response = rdf_flows.GrrMessage(
session_id=message.session_id,
name=message.name,
response_id=response_id,
request_id=message.request_id,
payload=response,
type=msg_type)
# Next expected response
response_id = response.response_id + 1
self.PushToStateQueue(manager, response)
# Status may only be None if the client reported itself as crashed.
if status is not None:
self.AddResourceUsage(status)
self.PushToStateQueue(
manager,
message,
response_id=response_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS)
else:
# Status may be None only if status_message_enforced is False.
if self.status_message_enforced:
raise RuntimeError(
"status message can only be None when "
"status_message_enforced is False")
# Additionally schedule a task for the worker
manager.QueueNotification(session_id=message.session_id,
priority=message.priority)
return len(request_tasks)
def __init__(self, request=None, responses=None):
self.status = None # A GrrStatus rdfvalue object.
self.success = True
self.request = request
if request:
self.request_data = rdf_protodict.Dict(request.data)
self._responses = []
self._dropped_responses = []
if responses:
# This may not be needed if we can assume that responses are
# returned in lexical order from the data_store.
responses.sort(key=operator.attrgetter("response_id"))
# The iterator that was returned as part of these responses. This should
# be passed back to actions that expect an iterator.
self.iterator = None
# Filter the responses by authorized states
for msg in responses:
# Check if the message is authenticated correctly.
if msg.auth_state != msg.AuthorizationState.AUTHENTICATED:
logging.warning(
"%s: Messages must be authenticated (Auth state %s)",
msg.session_id, msg.auth_state)
self._dropped_responses.append(msg)
# Skip this message - it is invalid
continue
# Check for iterators
if msg.type == msg.Type.ITERATOR:
self.iterator = rdf_client.Iterator(msg.payload)
continue
# Look for a status message
if msg.type == msg.Type.STATUS:
# Our status is set to the first status message that we see in
# the responses. We ignore all other messages after that.
self.status = rdf_flows.GrrStatus(msg.payload)
# Check this to see if the call succeeded
self.success = self.status.status == self.status.ReturnedStatus.OK
# Ignore all other messages
break
# Use this message
self._responses.append(msg)
if self.status is None:
# This is a special case of de-synchronized messages.
if self._dropped_responses:
logging.error(
"De-synchronized messages detected:\n %s", "\n".join([
utils.SmartUnicode(x)
for x in self._dropped_responses
]))
if responses:
self._LogFlowState(responses)
raise FlowError("No valid Status message.")
# This is the raw message accessible while going through the iterator
self.message = None
