Пример #1
0
 def StatFile(self, list_dir_req):
   if list_dir_req.pathspec.path == "/proc/kcore":
     result = rdf_client.StatEntry(pathspec=list_dir_req.pathspec,
                                   st_mode=400)
     status = rdf_flows.GrrStatus(status=rdf_flows.GrrStatus.ReturnedStatus.OK)
     return [result, status]
   raise IOError("Not found.")
Пример #2
0
 def RekallAction(self, rekall_request):
   if rekall_request.device.path != "/proc/kcore":
     return [rdf_flows.GrrStatus(
         status=rdf_flows.GrrStatus.ReturnedStatus.GENERIC_ERROR,
         error_message="Should use kcore device when present.")]
   response = rdf_rekall_types.RekallResponse(json_messages="{}")
   return [response, rdf_client.Iterator(state="FINISHED")]
Пример #3
0
  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, store=self.data_store) as manager:
      sessions_handled = []
      for session_id, msgs in utils.GroupBy(
          messages, operator.attrgetter("session_id")).iteritems():

        # Remove and handle messages to WellKnownFlows
        unprocessed_msgs = self.HandleWellKnownFlows(msgs)

        if not unprocessed_msgs:
          continue

        # Keep track of all the flows we handled in this request.
        sessions_handled.append(session_id)

        for msg in unprocessed_msgs:
          manager.QueueResponse(session_id, 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 use .Get() instead.
            if msg.HasTaskID():
              manager.DeQueueClientRequest(client_id, msg.task_id)

            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.
              events.Events.PublishEvent(
                  "ClientCrash", rdf_flows.GrrMessage(msg), token=self.token)

    logging.debug("Received %s messages in %s sec",
                  len(messages), time.time() - now)
Пример #4
0
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.client_worker.SendReply(
      session_id=rdf_flows.FlowSessionID(flow_name="InstallationFailed"),
      message_type=rdf_flows.GrrMessage.Type.STATUS,
      request_id=0,
      response_id=0,
      rdf_value=rdf_flows.GrrStatus(
          status=rdf_flows.GrrStatus.ReturnedStatus.GENERIC_ERROR,
          error_message="Installation failed.",
          backtrace=log_data[-10000:]))

  client.RunOnce()
Пример #5
0
    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)
Пример #6
0
  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)

    # 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 = rdf_flows.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
Пример #7
0
    def StatFile(self, args):
        """StatFile action mock."""
        req = rdf_client.ListDirRequest(args)

        response = rdf_client.StatEntry(pathspec=req.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)
        status.cpu_time_used.user_cpu_time = self.responses
        status.cpu_time_used.system_cpu_time = self.responses * 2
        status.network_bytes_sent = self.responses * 3

        # Every "failrate" client does not have this file.
        if self.count == self.failrate:
            self.count = 0
            return [status]

        return [response, status]
Пример #8
0
    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 = rdf_flows.GrrStatus(
                status=rdf_flows.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=rdf_flows.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)
Пример #9
0
    def SendResponse(self, session_id, data, client_id=None, well_known=False):
        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 = 1, 1
        with queue_manager.QueueManager(token=self.token) as flow_manager:
            flow_manager.QueueResponse(
                session_id,
                rdf_flows.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,
                    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,
                        type=rdf_flows.GrrMessage.Type.STATUS))

            flow_manager.QueueNotification(session_id=session_id)
            timestamp = flow_manager.frozen_timestamp

        return timestamp
Пример #10
0
  def testReceiveUnsolicitedClientMessage(self):
    client_id = test_lib.TEST_CLIENT_ID
    flow_obj = self.FlowSetup(
        flow_test_lib.FlowOrderTest.__name__, client_id=client_id)

    session_id = flow_obj.session_id
    status = rdf_flows.GrrStatus(status=rdf_flows.GrrStatus.ReturnedStatus.OK)
    messages = [
        # This message has no task_id set...
        rdf_flows.GrrMessage(
            request_id=1,
            response_id=1,
            session_id=session_id,
            payload=rdfvalue.RDFInteger(1),
            task_id=15),
        rdf_flows.GrrMessage(
            request_id=1,
            response_id=2,
            session_id=session_id,
            payload=status,
            type=rdf_flows.GrrMessage.Type.STATUS)
    ]

    self.server.ReceiveMessages(client_id, messages)
    manager = queue_manager.QueueManager(token=self.token)
    completed = list(manager.FetchCompletedRequests(session_id))
    self.assertEqual(len(completed), 1)
Пример #11
0
    def Error(self, backtrace, client_id=None, status=None):
        """Kills this flow with an error."""
        client_id = client_id or self.runner_args.client_id
        if self.IsRunning():
            # Set an error status
            reply = rdf_flows.GrrStatus()
            if status is None:
                reply.status = rdf_flows.GrrStatus.ReturnedStatus.GENERIC_ERROR
            else:
                reply.status = status

            if backtrace:
                reply.error_message = backtrace

            self.flow_obj.Terminate(status=reply)

            self.context.state = rdf_flows.FlowContext.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)
Пример #12
0
    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)
Пример #13
0
    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_flows.FlowContext.State.ERROR
        self.Notify("FlowStatus", client_id,
                    "Flow (%s) terminated due to error" % self.session_id)

        self.flow_obj.Flush()
Пример #14
0
    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 != rdf_flows.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 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
            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 = 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=response)

            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 = rdf_flows.Flow.State.TERMINATED
        self.flow_obj.Flush()
Пример #15
0
 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(
         rdf_flows.GrrStatus(),
         session_id=rdfvalue.SessionID("W:session"),
         response_id=i, request_id=1)
Пример #16
0
    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 = flows.GrrStatus(status=flows.GrrStatus.ReturnedStatus.OK)
Пример #17
0
  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
Пример #18
0
  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)
Пример #19
0
  def testEqualTimestampNotifications(self):
    frontend_server = front_end.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.GRRFlow.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)
Пример #20
0
  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
Пример #21
0
  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_lib.CONFIG["Client.gc_frequency"]
Пример #22
0
    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()
Пример #23
0
    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 != rdf_flows.FlowContext.State.RUNNING:
            return

        if self.runner_args.request_state.session_id:
            # Make a response or use the existing one.
            response = status or 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
            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.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=response)

            try:
                # Queue the response now
                self.queue_manager.QueueResponse(msg)
            finally:
                self.QueueNotification(session_id=request_state.session_id)

        # Mark as terminated.
        self.context.state = rdf_flows.FlowContext.State.TERMINATED
        self.flow_obj.Flush()
Пример #24
0
  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)

    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)
Пример #25
0
  def testReceiveMessagesWithStatus(self):
    """Receiving a sequence of messages with a status."""
    client_id = test_lib.TEST_CLIENT_ID
    flow_obj = self.FlowSetup(
        flow_test_lib.FlowOrderTest.__name__, client_id=client_id)

    session_id = flow_obj.session_id
    messages = [
        rdf_flows.GrrMessage(
            request_id=1,
            response_id=i,
            session_id=session_id,
            payload=rdfvalue.RDFInteger(i),
            task_id=15) for i in range(1, 10)
    ]

    # Now add the status message
    status = rdf_flows.GrrStatus(status=rdf_flows.GrrStatus.ReturnedStatus.OK)
    messages.append(
        rdf_flows.GrrMessage(
            request_id=1,
            response_id=len(messages) + 1,
            task_id=15,
            session_id=messages[0].session_id,
            payload=status,
            type=rdf_flows.GrrMessage.Type.STATUS))

    self.server.ReceiveMessages(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(client_id, 100)
    self.assertEqual(len(tasks_on_client_queue), 1)

    stored_messages = data_store.DB.ReadResponsesForRequestId(session_id, 1)

    self.assertEqual(len(stored_messages), len(messages))

    stored_messages.sort(key=lambda m: m.response_id)
    # Check that messages were stored correctly
    for stored_message, message in zip(stored_messages, messages):
      # We don't care about the last queueing time.
      stored_message.timestamp = None
      self.assertRDFValuesEqual(stored_message, message)
Пример #26
0
  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(
            rdf_flows.GrrStatus(
                status=rdf_flows.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=rdf_flows.GrrMessage.Type.STATUS)
        if flags.FLAGS.debug:
          pdb.post_mortem()
Пример #27
0
 def _ScheduleResponseAndStatus(self, client_id, flow_id):
   with queue_manager.QueueManager(token=self.token) as flow_manager:
     # Schedule a response.
     flow_manager.QueueResponse(
         rdf_flows.GrrMessage(
             source=client_id,
             session_id=flow_id,
             payload=rdf_protodict.DataBlob(string="Helllo"),
             request_id=1,
             response_id=1))
     # And a STATUS message.
     flow_manager.QueueResponse(
         rdf_flows.GrrMessage(
             source=client_id,
             session_id=flow_id,
             payload=rdf_flows.GrrStatus(
                 status=rdf_flows.GrrStatus.ReturnedStatus.OK),
             request_id=1,
             response_id=2,
             type=rdf_flows.GrrMessage.Type.STATUS))
Пример #28
0
    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 = rdf_flows.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 = rdf_flows.GrrMessage(
                    session_id=hunt_id,
                    request_id=state.id,
                    response_id=1,
                    auth_state=rdf_flows.GrrMessage.AuthorizationState.
                    AUTHENTICATED,
                    type=rdf_flows.GrrMessage.Type.STATUS,
                    payload=rdf_flows.GrrStatus())

                flow_manager.QueueResponse(hunt_id, msg)

                # And notify the worker about it.
                flow_manager.QueueNotification(session_id=hunt_id)
Пример #29
0
    def testReceiveMessagesWithStatus(self):
        """Receiving a sequence of messages with a status."""
        flow_obj = self.FlowSetup("FlowOrderTest")

        session_id = flow_obj.session_id
        messages = [
            rdf_flows.GrrMessage(request_id=1,
                                 response_id=i,
                                 session_id=session_id,
                                 payload=rdfvalue.RDFInteger(i),
                                 task_id=15) for i in range(1, 10)
        ]

        # Now add the status message
        status = rdf_flows.GrrStatus(
            status=rdf_flows.GrrStatus.ReturnedStatus.OK)
        messages.append(
            rdf_flows.GrrMessage(request_id=1,
                                 response_id=len(messages) + 1,
                                 task_id=15,
                                 session_id=messages[0].session_id,
                                 payload=status,
                                 type=rdf_flows.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 = rdf_flows.GrrMessage.FromSerializedString(
                stored_message)
            self.assertRDFValuesEqual(stored_message, message)
Пример #30
0
    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(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,
                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,
            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.

        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.context.backtrace)
        self.assertNotEqual(flow_obj.context.state,
                            rdf_flows.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"])