def OnStartup(self): """A handler that is called on client startup.""" # We read the transaction log and fail any requests that are in it. If there # is anything in the transaction log we assume its there because we crashed # last time and let the server know. last_request = self.nanny_controller.GetTransactionLog() if last_request: status = rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.CLIENT_KILLED, error_message="Client killed during transaction") nanny_status = self.nanny_controller.GetNannyStatus() if nanny_status: status.nanny_status = nanny_status self.SendReply(status, request_id=last_request.request_id, response_id=1, session_id=last_request.session_id, message_type=rdfvalue.GrrMessage.Type.STATUS) self.nanny_controller.CleanTransactionLog() # Inform the server that we started. action_cls = actions.ActionPlugin.classes.get("SendStartupInfo", actions.ActionPlugin) action = action_cls(grr_worker=self) action.Run(None, ttl=1)
def testReceiveMessagesWithStatus(self): """Receiving a sequence of messages with a status.""" flow_obj = self.FlowSetup("FlowOrderTest") session_id = flow_obj.session_id messages = [rdfvalue.GrrMessage(request_id=1, response_id=i, session_id=session_id, args=str(i), task_id=15) for i in range(1, 10)] # Now add the status message status = rdfvalue.GrrStatus(status=rdfvalue.GrrStatus.ReturnedStatus.OK) messages.append(rdfvalue.GrrMessage( request_id=1, response_id=len(messages)+1, task_id=15, session_id=messages[0].session_id, payload=status, type=rdfvalue.GrrMessage.Type.STATUS)) self.server.ReceiveMessages(self.client_id, messages) # Make sure the task is still on the client queue manager = queue_manager.QueueManager(token=self.token) tasks_on_client_queue = manager.Query(self.client_id, 100) self.assertEqual(len(tasks_on_client_queue), 1) # Check that messages were stored correctly for message in messages: stored_message, _ = data_store.DB.Resolve( session_id.Add("state/request:00000001"), manager.FLOW_RESPONSE_TEMPLATE % (1, message.response_id), token=self.token) stored_message = rdfvalue.GrrMessage(stored_message) self.assertProtoEqual(stored_message, message)
def run(self): """Main thread for processing messages.""" self.OnStartup() # As long as our output queue has some room we can process some # input messages: while True: message = self._in_queue.get() # A message of None is our terminal message. if message is None: break try: self.HandleMessage(message) # Catch any errors and keep going here except Exception as e: # pylint: disable=broad-except logging.warn("%s", e) self.SendReply(rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.GENERIC_ERROR, error_message=utils.SmartUnicode(e)), request_id=message.request_id, response_id=message.response_id, session_id=message.session_id, task_id=message.task_id, message_type=rdfvalue.GrrMessage.Type.STATUS) if flags.FLAGS.debug: pdb.post_mortem()
def InstallerNotifyServer(): """An emergency function Invoked when the client installation failed.""" # We make a temporary emergency config file to contain the new client id. Note # that the notification callback does not really mean anything to us, since # the client is not installed and we dont have basic interrogate information. config_lib.CONFIG.SetWriteBack("temp.yaml") try: log_data = open(config_lib.CONFIG["Installer.logfile"], "rb").read() except (IOError, OSError): log_data = "" # Start the client and send the server a message, then terminate. The # private key may be empty if we did not install properly yet. In this case, # the client will automatically generate a random client ID and private key # (and the message will be unauthenticated since we never enrolled.). comms.CommsInit().RunOnce() client = comms.GRRHTTPClient( ca_cert=config_lib.CONFIG["CA.certificate"], private_key=config_lib.CONFIG.Get("Client.private_key")) client.EstablishConnection() client.client_worker.SendReply( session_id="W:InstallationFailed", message_type=rdfvalue.GrrMessage.Type.STATUS, request_id=0, response_id=0, rdf_value=rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.GENERIC_ERROR, error_message="Installation failed.", backtrace=log_data[-10000:])) client.RunOnce()
def StartFlow(self, client_id): flow_id = flow.GRRFlow.StartFlow(client_id=client_id, flow_name="RecursiveListDirectory", max_depth=5, queue=self.queue, token=self.token) self.flow_ids.append(flow_id) messages = [] for d in range(self.nr_dirs): messages += self.GenerateFiles(client_id, self.files_per_dir, "dir/dir%d" % d) messages.append(rdfvalue.GrrStatus()) with queue_manager.QueueManager(token=self.token) as flow_manager: for i, payload in enumerate(messages): msg = rdfvalue.GrrMessage(session_id=flow_id, request_id=1, response_id=1 + i, auth_state=rdfvalue.GrrMessage. AuthorizationState.AUTHENTICATED, payload=payload) if isinstance(payload, rdfvalue.GrrStatus): msg.type = 1 flow_manager.QueueResponse(flow_id, msg)
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 = rdfvalue.GrrMessage( name="MockAction", session_id=self.session_id, auth_state=rdfvalue.GrrMessage.AuthorizationState.UNAUTHENTICATED, request_id=1) 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 = rdfvalue.GrrStatus(message_list[0].args) self.assert_("not Authenticated" in status.error_message) self.assert_("RuntimeError" in status.error_message) self.assertNotEqual(status.status, rdfvalue.GrrStatus.ReturnedStatus.OK)
def Error(self, backtrace, client_id=None, status=None): """Kills this flow with an error.""" client_id = client_id or self.args.client_id if self.IsRunning(): # Set an error status reply = rdfvalue.GrrStatus() if status is None: reply.status = rdfvalue.GrrStatus.ReturnedStatus.GENERIC_ERROR else: reply.status = status if backtrace: reply.error_message = backtrace self.Terminate(status=reply) self.context.state = rdfvalue.Flow.State.ERROR if 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.Notify("FlowStatus", client_id, "Flow (%s) terminated due to error" % self.session_id)
def SendOKStatus(self, response_id, session_id): """Send a message to the flow.""" message = rdfvalue.GrrMessage( request_id=1, response_id=response_id, session_id=session_id, type=rdfvalue.GrrMessage.Type.STATUS, auth_state=rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED) status = rdfvalue.GrrStatus(status=rdfvalue.GrrStatus.ReturnedStatus.OK) message.payload = status self.SendMessage(message) # Now also set the state on the RequestState request_state, _ = data_store.DB.Resolve( message.session_id.Add("state"), queue_manager.QueueManager.FLOW_REQUEST_TEMPLATE % message.request_id, token=self.token) request_state = rdfvalue.RequestState(request_state) request_state.status = status data_store.DB.Set( message.session_id.Add("state"), queue_manager.QueueManager.FLOW_REQUEST_TEMPLATE % message.request_id, request_state, token=self.token) return message
def SendResponse(self, session_id, data, client_id=None, well_known=False): if not isinstance(data, rdfvalue.RDFValue): data = rdfvalue.DataBlob(string=data) if well_known: request_id, response_id = 0, 12345 else: request_id, response_id = 1, 1 with queue_manager.QueueManager(token=self.token) as flow_manager: flow_manager.QueueResponse( session_id, rdfvalue.GrrMessage(source=client_id, session_id=session_id, payload=data, request_id=request_id, response_id=response_id)) if not well_known: # For normal flows we have to send a status as well. flow_manager.QueueResponse( session_id, rdfvalue.GrrMessage( source=client_id, session_id=session_id, payload=rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.OK), request_id=request_id, response_id=response_id + 1, type=rdfvalue.GrrMessage.Type.STATUS)) # Signal on the worker queue that this flow is ready. data_store.DB.Set(worker.DEFAULT_WORKER_QUEUE, "task:%s" % session_id, "X", token=self.token)
def SendResponse(self, session_id, data, client_id=None, well_known=False): if not isinstance(data, rdfvalue.RDFValue): data = rdfvalue.DataBlob(string=data) if well_known: request_id, response_id = 0, 12345 else: request_id, response_id = 1, 1 with queue_manager.QueueManager(token=self.token) as flow_manager: flow_manager.QueueResponse( session_id, rdfvalue.GrrMessage(source=client_id, session_id=session_id, payload=data, request_id=request_id, response_id=response_id)) if not well_known: # For normal flows we have to send a status as well. flow_manager.QueueResponse( session_id, rdfvalue.GrrMessage( source=client_id, session_id=session_id, payload=rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.OK), request_id=request_id, response_id=response_id + 1, type=rdfvalue.GrrMessage.Type.STATUS)) flow_manager.QueueNotification(session_id=session_id) timestamp = flow_manager.frozen_timestamp return timestamp
def StatFile(self, list_dir_req): if list_dir_req.pathspec.path == "/proc/kcore": result = rdfvalue.StatEntry(pathspec=list_dir_req.pathspec, st_mode=400) status = rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.OK) return [result, status] raise IOError("Not found.")
def Terminate(self, status=None): """Terminates this flow.""" try: self.queue_manager.DestroyFlowStates(self.session_id) except queue_manager.MoreDataException: pass # This flow might already not be running. if self.context.state != rdfvalue.Flow.State.RUNNING: return try: # Close off the output collection. if self.output and len(self.output): self.output.Close() logging.info("%s flow results written to %s", len(self.output), self.output.urn) self.output = None except access_control.UnauthorizedAccess: # This might fail if the output has a pickled token. pass if self.args.request_state.session_id: logging.debug("Terminating flow %s", self.session_id) # Make a response or use the existing one. response = status or rdfvalue.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 response.cpu_time_used.user_cpu_time = user_cpu response.cpu_time_used.system_cpu_time = sys_cpu response.network_bytes_sent = self.context.network_bytes_sent response.child_session_id = self.session_id request_state = self.args.request_state request_state.response_count += 1 # Make a response message msg = rdfvalue.GrrMessage(session_id=request_state.session_id, request_id=request_state.id, response_id=request_state.response_count, auth_state=rdfvalue.GrrMessage. AuthorizationState.AUTHENTICATED, type=rdfvalue.GrrMessage.Type.STATUS, args=response.SerializeToString()) try: # Queue the response now self.queue_manager.QueueResponse(request_state.session_id, msg) finally: self.QueueNotification(session_id=request_state.session_id) # Mark as terminated. self.context.state = rdfvalue.Flow.State.TERMINATED self.flow_obj.Flush()
def RekallAction(self, rekall_request): if rekall_request.device.path != "/proc/kcore": return [ rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.GENERIC_ERROR, error_message="Should use kcore device when present.") ] response = rdfvalue.RekallResponse(json_messages="{}") return [response, rdfvalue.Iterator(state="FINISHED")]
def SendToServer(self): """Schedule some packets from client to server.""" # Generate some client traffic for i in range(0, 10): self.client_communicator.client_worker.SendReply( rdfvalue.GrrStatus(), session_id=rdfvalue.SessionID("W:session"), response_id=i, request_id=1)
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 = rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.OK)
def Run(self, unused_arg): """Run the kill.""" # Send a message back to the service to say that we are about to shutdown. reply = rdfvalue.GrrStatus(status=rdfvalue.GrrStatus.ReturnedStatus.OK) # Queue up the response message, jump the queue. self.SendReply(reply, message_type=rdfvalue.GrrMessage.Type.STATUS, priority=rdfvalue.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 _ScheduleResponseAndStatus(self, client_id, flow_id): with queue_manager.QueueManager(token=self.token) as flow_manager: # Schedule a response. flow_manager.QueueResponse(flow_id, rdfvalue.GrrMessage( source=client_id, session_id=flow_id, payload=rdfvalue.DataBlob(string="Helllo"), request_id=1, response_id=1)) # And a STATUS message. flow_manager.QueueResponse(flow_id, rdfvalue.GrrMessage( source=client_id, session_id=flow_id, payload=rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.OK), request_id=1, response_id=2, type=rdfvalue.GrrMessage.Type.STATUS))
def testHandleError(self): """Test handling of a request which raises.""" # Push a request on it message = rdfvalue.GrrMessage( name="RaiseAction", session_id=self.session_id, auth_state=rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED, request_id=1) 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 = rdfvalue.GrrStatus(message_list[0].args) self.assert_("RuntimeError" in status.error_message) self.assertNotEqual(status.status, rdfvalue.GrrStatus.ReturnedStatus.OK)
def QueueMessages(self, messages): """Queue a message from the server for processing. We maintain all the incoming messages in a queue. These messages are consumed until the outgoing queue fills to the allowable level. This mechanism allows us to throttle the server messages and limit the size of the outgoing queue on the client. Note that we can only limit processing of single request messages so if a single request message generates huge amounts of response messages we will still overflow the output queue. Therefore actions must be written in such a way that each request generates a limited and known maximum number and size of responses. (e.g. do not write a single client action to fetch the entire disk). Args: messages: List of parsed protobuf arriving from the server. """ # Push all the messages to our input queue for message in messages: self._in_queue.append(message) stats.STATS.IncrementCounter("grr_client_received_messages") # As long as our output queue has some room we can process some # input messages: while self._in_queue and ( self._out_queue_size < config_lib.CONFIG["Client.max_out_queue"]): message = self._in_queue.pop(0) try: self.HandleMessage(message) # Catch any errors and keep going here except Exception as e: # pylint: disable=broad-except self.SendReply( rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.GENERIC_ERROR, error_message=utils.SmartUnicode(e)), request_id=message.request_id, response_id=message.response_id, session_id=message.session_id, task_id=message.task_id, message_type=rdfvalue.GrrMessage.Type.STATUS) if flags.FLAGS.debug: pdb.post_mortem()
def StartClients(cls, hunt_id, client_ids, token=None): """This method is called by the foreman for each client it discovers. Note that this function is performance sensitive since it is called by the foreman for every client which needs to be scheduled. Args: hunt_id: The hunt to schedule. client_ids: List of clients that should be added to the hunt. token: An optional access token to use. """ token = token or access_control.ACLToken(username="******", reason="hunting") with queue_manager.QueueManager(token=token) as flow_manager: for client_id in client_ids: # Now we construct a special response which will be sent to the hunt # flow. Randomize the request_id so we do not overwrite other messages # in the queue. state = rdfvalue.RequestState(id=utils.PRNG.GetULong(), session_id=hunt_id, client_id=client_id, next_state="AddClient") # Queue the new request. flow_manager.QueueRequest(hunt_id, state) # Send a response. msg = rdfvalue.GrrMessage(session_id=hunt_id, request_id=state.id, response_id=1, auth_state=rdfvalue.GrrMessage. AuthorizationState.AUTHENTICATED, type=rdfvalue.GrrMessage.Type.STATUS, payload=rdfvalue.GrrStatus()) flow_manager.QueueResponse(hunt_id, msg) # And notify the worker about it. flow_manager.QueueNotification(session_id=hunt_id)
def CallState(self, messages=None, next_state="", client_id=None, request_data=None, start_time=None): """This method is used to asynchronously schedule a new hunt state. 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 hunt to be invoked with the responses. client_id: ClientURN to use in scheduled requests. 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: Schedule the state at this time. This delays notification and messages for processing into the future. Raises: ValueError: on arguments error. """ if messages is None: messages = [] if not next_state: raise ValueError("next_state can't be empty.") # Now we construct a special response which will be sent to the hunt # flow. Randomize the request_id so we do not overwrite other messages in # the queue. request_state = rdfvalue.RequestState( id=utils.PRNG.GetULong(), session_id=self.context.session_id, client_id=client_id, next_state=next_state) if request_data: request_state.data = rdfvalue.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], rdfvalue.GrrStatus): messages.append(rdfvalue.GrrStatus()) # Send all the messages for i, payload in enumerate(messages): if isinstance(payload, rdfvalue.RDFValue): msg = rdfvalue.GrrMessage( session_id=self.session_id, request_id=request_state.id, response_id=1 + i, auth_state=rdfvalue.GrrMessage.AuthorizationState. AUTHENTICATED, payload=payload, type=rdfvalue.GrrMessage.Type.MESSAGE) if isinstance(payload, rdfvalue.GrrStatus): msg.type = rdfvalue.GrrMessage.Type.STATUS else: raise flow_runner.FlowRunnerError( "Bad message %s of type %s." % (payload, type(payload))) self.QueueResponse(msg, timestamp=start_time) # Add the status message if needed. if not messages or not isinstance(messages[-1], rdfvalue.GrrStatus): messages.append(rdfvalue.GrrStatus()) # Notify the worker about it. self.QueueNotification(session_id=self.session_id, timestamp=start_time)
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 = rdfvalue.RequestState( id=self.GetNextOutboundId(), session_id=self.context.session_id, client_id=self.args.client_id, next_state=next_state) if request_data: request_state.data = rdfvalue.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], rdfvalue.GrrStatus): messages.append(rdfvalue.GrrStatus()) # Send all the messages for i, payload in enumerate(messages): if isinstance(payload, rdfvalue.RDFValue): msg = rdfvalue.GrrMessage( session_id=self.session_id, request_id=request_state.id, response_id=1 + i, auth_state=rdfvalue.GrrMessage.AuthorizationState. AUTHENTICATED, payload=payload, type=rdfvalue.GrrMessage.Type.MESSAGE) if isinstance(payload, rdfvalue.GrrStatus): msg.type = rdfvalue.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 Execute(self): """This function parses the RDFValue from the server. The Run method will be called with the specified RDFValue. Returns: Upon return a callback will be called on the server to register the end of the function and pass back exceptions. Raises: RuntimeError: The arguments from the server do not match the expected rdf type. """ args = None try: if self.message.args_rdf_name: if not self.in_rdfvalue: raise RuntimeError("Did not expect arguments, got %s." % self.message.args_rdf_name) if self.in_rdfvalue.__name__ != self.message.args_rdf_name: raise RuntimeError("Unexpected arg type %s != %s." % (self.message.args_rdf_name, self.in_rdfvalue.__name__)) args = self.message.payload self.status = rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.OK) # Only allow authenticated messages in the client if (self._authentication_required and self.message.auth_state != rdfvalue.GrrMessage.AuthorizationState.AUTHENTICATED): raise RuntimeError("Message for %s was not Authenticated." % self.message.name) pid = os.getpid() self.proc = psutil.Process(pid) user_start, system_start = self.proc.get_cpu_times() self.cpu_start = (user_start, system_start) self.cpu_limit = self.message.cpu_limit self.network_bytes_limit = self.message.network_bytes_limit try: self.Run(args) # Ensure we always add CPU usage even if an exception occured. finally: user_end, system_end = self.proc.get_cpu_times() self.cpu_used = (user_end - user_start, system_end - system_start) except NetworkBytesExceededError as e: self.SetStatus( rdfvalue.GrrStatus.ReturnedStatus.NETWORK_LIMIT_EXCEEDED, "%r: %s" % (e, e), traceback.format_exc()) # We want to report back all errors and map Python exceptions to # Grr Errors. except Exception as e: # pylint: disable=broad-except self.SetStatus(rdfvalue.GrrStatus.ReturnedStatus.GENERIC_ERROR, "%r: %s" % (e, e), traceback.format_exc()) if flags.FLAGS.debug: pdb.post_mortem() if self.status.status != rdfvalue.GrrStatus.ReturnedStatus.OK: logging.info("Job Error (%s): %s", self.__class__.__name__, self.status.error_message) if self.status.backtrace: logging.debug(self.status.backtrace) if self.cpu_used: self.status.cpu_time_used.user_cpu_time = self.cpu_used[0] self.status.cpu_time_used.system_cpu_time = self.cpu_used[1] # This returns the error status of the Actions to the flow. self.SendReply(self.status, message_type=rdfvalue.GrrMessage.Type.STATUS)
def ProcessMessage(self, message=None, event=None): """Processes this event.""" _ = event client_id = message.source nanny_msg = "" flow_obj = aff4.FACTORY.Open(message.session_id, token=self.token) # Log. logging.info("Client crash reported, client %s.", client_id) # Export. stats.STATS.IncrementCounter("grr_client_crashes") # Write crash data to AFF4. client = aff4.FACTORY.Open(client_id, token=self.token) client_info = client.Get(client.Schema.CLIENT_INFO) status = rdfvalue.GrrStatus(message.args) crash_details = rdfvalue.ClientCrash( client_id=client_id, session_id=message.session_id, client_info=client_info, crash_message=status.error_message, timestamp=rdfvalue.RDFDatetime().Now(), crash_type=self.well_known_session_id) self.WriteAllCrashDetails(client_id, crash_details, flow_session_id=message.session_id) # Also send email. if config_lib.CONFIG["Monitoring.alert_email"]: if status.nanny_status: nanny_msg = "Nanny status: %s" % status.nanny_status client = aff4.FACTORY.Open(client_id, token=self.token) hostname = client.Get(client.Schema.HOSTNAME) url = urllib.urlencode( (("c", client_id), ("main", "HostInformation"))) renderer = rendering.FindRendererForObject(flow_obj.state) email_alerts.SendEmail( config_lib.CONFIG["Monitoring.alert_email"], "GRR server", "Client %s reported a crash." % client_id, self.mail_template % dict(client_id=client_id, admin_ui=config_lib.CONFIG["AdminUI.url"], hostname=hostname, state=renderer.RawHTML(), urn=url, nanny_msg=nanny_msg, signature=config_lib.CONFIG["Email.signature"]), is_html=True) if nanny_msg: msg = "Client crashed, " + nanny_msg else: msg = "Client crashed." # Now terminate the flow. flow.GRRFlow.TerminateFlow(message.session_id, reason=msg, token=self.token, force=True)
def testEnums(self): """Check that enums are wrapped in a descriptor class.""" sample = rdfvalue.GrrStatus() self.assertEqual(str(sample.status), "OK")
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.GRRFlow.StartFlow(client_id=self.client_id, flow_name="WorkerSendingTestFlow", token=self.token) worker_obj = worker.GRRWorker(worker.DEFAULT_WORKER_QUEUE, 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( session_id, rdfvalue.GrrMessage(source=self.client_id, session_id=session_id, payload=rdfvalue.DataBlob(string="Response 2"), request_id=request_id, response_id=response_id)) status = rdfvalue.GrrMessage( source=self.client_id, session_id=session_id, payload=rdfvalue.GrrStatus( status=rdfvalue.GrrStatus.ReturnedStatus.OK), request_id=request_id, response_id=response_id + 1, type=rdfvalue.GrrMessage.Type.STATUS) # Now we write half the status information. subject = session_id.Add("state") queue = flow_manager.to_write.setdefault(subject, {}) queue.setdefault(flow_manager.FLOW_STATUS_TEMPLATE % request_id, []).append((status.SerializeToString(), None)) flow_manager.Flush() # 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. 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) 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.state.context.backtrace) self.assertNotEqual(flow_obj.state.context.state, rdfvalue.Flow.State.ERROR) request2_data = data_store.DB.ResolveRegex(session_id.Add("state"), ".*:00000002", token=self.token) # Make sure the status field and the original request are still there. self.assertEqual(len(request2_data), 2) request1_data = data_store.DB.ResolveRegex(session_id.Add("state"), ".*:00000001", token=self.token) # Everything from request 1 should have been deleted. self.assertEqual(len(request1_data), 0) # The notification for request 2 should have survived. with queue_manager.QueueManager(token=self.token) as manager: notifications = manager.GetNotifications( worker.DEFAULT_WORKER_QUEUE) 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(session_id, 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"])