def testForemanMessageHandler(self): with mock.patch.object(foreman.Foreman, "AssignTasksToClient") as instr: # Send a message to the Foreman. client_id = "C.1100110011001100" data_store.REL_DB.WriteMessageHandlerRequests([ rdf_objects.MessageHandlerRequest( client_id=client_id, handler_name="ForemanHandler", request_id=12345, request=rdf_protodict.DataBlob()) ]) done = threading.Event() def handle(l): worker_lib.ProcessMessageHandlerRequests(l) done.set() data_store.REL_DB.RegisterMessageHandler( handle, worker_lib.GRRWorker.message_handler_lease_time, limit=1000) try: self.assertTrue(done.wait(10)) # Make sure there are no leftover requests. self.assertEqual( data_store.REL_DB.ReadMessageHandlerRequests(), []) instr.assert_called_once_with(client_id) finally: data_store.REL_DB.UnregisterMessageHandler(timeout=60)
def testMessageHandlerRequests(self): requests = [ rdf_objects.MessageHandlerRequest(client_id="C.1000000000000000", handler_name="Testhandler", request_id=i * 100, request=rdfvalue.RDFInteger(i)) for i in range(5) ] self.db.WriteMessageHandlerRequests(requests) read = self.db.ReadMessageHandlerRequests() for r in read: self.assertTrue(r.timestamp) r.timestamp = None self.assertEqual(sorted(read, key=lambda req: req.request_id), requests) self.db.DeleteMessageHandlerRequests(requests[:2]) self.db.DeleteMessageHandlerRequests(requests[4:5]) read = self.db.ReadMessageHandlerRequests() self.assertLen(read, 2) read = sorted(read, key=lambda req: req.request_id) for r in read: r.timestamp = None self.assertEqual(requests[2:4], read) self.db.DeleteMessageHandlerRequests(read)
def testEnrollmentHandler(self): self._ClearClient() # First 406 queues an EnrolmentRequest. status = self.client_communicator.RunOnce() self.assertEqual(status.code, 406) # Send it to the server. status = self.client_communicator.RunOnce() self.assertEqual(status.code, 406) self.assertLen(self.messages, 1) self.assertEqual(self.messages[0].session_id.Basename(), "E:%s" % ca_enroller.EnrolmentHandler.handler_name) request = rdf_objects.MessageHandlerRequest( client_id=self.messages[0].source.Basename(), handler_name="Enrol", request_id=12345, request=self.messages[0].payload) handler = ca_enroller.EnrolmentHandler(token=self.token) handler.ProcessMessages([request]) # The next client communication should give a 200. status = self.client_communicator.RunOnce() self.assertEqual(status.code, 200)
def _PushHandlerMessage(self, message): """Pushes a message that goes to a message handler.""" # We only accept messages of type MESSAGE. if message.type != rdf_flows.GrrMessage.Type.MESSAGE: raise ValueError("Unexpected message type: %s" % type(message)) if not message.session_id: raise ValueError("Message without session_id: %s" % message) # Assume the message is authenticated and comes from this client. message.source = self.client_id message.auth_state = "AUTHENTICATED" session_id = message.session_id handler_name = message_handlers.session_id_map.get(session_id, None) if handler_name is None: raise ValueError("Unknown well known session id in msg %s" % message) logging.info("Running message handler: %s", handler_name) handler_cls = handler_registry.handler_name_map.get(handler_name) handler_request = rdf_objects.MessageHandlerRequest( client_id=self.client_id, handler_name=handler_name, request_id=message.response_id, request=message.payload) handler_cls().ProcessMessages([handler_request])
def _PushHandlerMessage(self, message): """Pushes a message that goes to a message handler.""" # We only accept messages of type MESSAGE. if message.type != rdf_flows.GrrMessage.Type.MESSAGE: raise ValueError("Unexpected message type: %s" % type(message)) if not message.session_id: raise ValueError("Message without session_id: %s" % message) # Assume the message is authenticated and comes from this client. message.source = self.client_id message.auth_state = "AUTHENTICATED" session_id = message.session_id if data_store.RelationalDBEnabled(): handler_name = queue_manager.session_id_map.get(session_id, None) if handler_name is None: raise ValueError("Unknown well known session id in msg %s" % message) logging.info("Running message handler: %s", handler_name) handler_cls = handler_registry.handler_name_map.get(handler_name) handler_request = rdf_objects.MessageHandlerRequest( client_id=self.client_id.Basename(), handler_name=handler_name, request_id=message.response_id, request=message.payload) handler_cls(token=self.token).ProcessMessages([handler_request]) else: logging.info("Running well known flow: %s", session_id) self.well_known_flows[session_id.FlowName()].ProcessMessage( message)
def _testProcessMessagesWellKnown(self): worker_obj = self._TestWorker() # Send a message to a WellKnownFlow - ClientStatsAuto. session_id = administrative.GetClientStatsAuto.well_known_session_id client_id = self.SetupClient(100) if data_store.RelationalDBReadEnabled(): done = threading.Event() def handle(l): worker_obj._ProcessMessageHandlerRequests(l) done.set() data_store.REL_DB.RegisterMessageHandler( handle, worker_obj.well_known_flow_lease_time, limit=1000) data_store.REL_DB.WriteMessageHandlerRequests([ rdf_objects.MessageHandlerRequest( client_id=client_id.Basename(), handler_name="StatsHandler", request_id=12345, request=rdf_client_stats.ClientStats(RSS_size=1234)) ]) self.assertTrue(done.wait(10)) else: self.SendResponse(session_id, data=rdf_client_stats.ClientStats(RSS_size=1234), client_id=client_id, well_known=True) # Process all messages worker_obj.RunOnce() worker_obj.thread_pool.Join() if data_store.RelationalDBReadEnabled(): results = data_store.REL_DB.ReadClientStats( client_id=client_id.Basename(), min_timestamp=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(0), max_timestamp=rdfvalue.RDFDatetime.Now()) self.assertLen(results, 1) stats = results[0] else: client = aff4.FACTORY.Open(client_id.Add("stats"), token=self.token) stats = client.Get(client.Schema.STATS) self.assertEqual(stats.RSS_size, 1234) # Make sure no notifications have been sent. user = aff4.FACTORY.Open("aff4:/users/%s" % self.token.username, token=self.token) notifications = user.Get(user.Schema.PENDING_NOTIFICATIONS) self.assertIsNone(notifications) if data_store.RelationalDBReadEnabled(): data_store.REL_DB.UnregisterMessageHandler(timeout=60)
def testEnrollment(self): """Test the http response to unknown clients.""" self._ClearClient() # Now communicate with the server. self.SendToServer() status = self.client_communicator.RunOnce() # We expect to receive a 406 and all client messages will be tagged as # UNAUTHENTICATED. self.assertEqual(status.code, 406) self.assertLen(self.messages, 10) self.assertEqual( self.messages[0].auth_state, rdf_flows.GrrMessage.AuthorizationState.UNAUTHENTICATED) # The next request should be an enrolling request. self.client_communicator.RunOnce() self.assertLen(self.messages, 11) enrolment_messages = [] expected_id = "E:%s" % ca_enroller.EnrolmentHandler.handler_name for m in self.messages: if m.session_id.Basename() == expected_id: enrolment_messages.append(m) self.assertLen(enrolment_messages, 1) # Now we manually run the enroll well known flow with the enrollment # request. This will start a new flow for enrolling the client, sign the # cert and add it to the data store. handler = ca_enroller.EnrolmentHandler() req = rdf_objects.MessageHandlerRequest( client_id=self.client_id, request=enrolment_messages[0].payload) handler.ProcessMessages([req]) # The next client communication should be enrolled now. status = self.client_communicator.RunOnce() self.assertEqual(status.code, 200) # There should be a cert for the client right now. md = data_store.REL_DB.ReadClientMetadata(self.client_id) self.assertTrue(md.certificate) # Now communicate with the server once again. self.SendToServer() status = self.client_communicator.RunOnce() self.assertEqual(status.code, 200)
def PushToStateQueue(self, manager, message, **kw): """Push given message to the state queue.""" # Assume the client is authorized message.auth_state = rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED # Update kw args for k, v in kw.items(): setattr(message, k, v) # Handle well known flows if message.request_id == 0: # Well known flows only accept messages of type MESSAGE. if message.type == rdf_flows.GrrMessage.Type.MESSAGE: # Assume the message is authenticated and comes from this client. message.source = self.client_id message.auth_state = "AUTHENTICATED" session_id = message.session_id if session_id: handler_name = queue_manager.session_id_map.get(session_id) if handler_name: logging.info("Running message handler: %s", handler_name) handler_cls = handler_registry.handler_name_map.get( handler_name) handler_request = rdf_objects.MessageHandlerRequest( client_id=self.client_id.Basename(), handler_name=handler_name, request_id=message.response_id, request=message.payload) handler_cls( token=self.token).ProcessMessages(handler_request) else: logging.info("Running well known flow: %s", session_id) self.well_known_flows[ session_id.FlowName()].ProcessMessage(message) return manager.QueueResponse(message)
def testMessageHandlerRequestLeasing(self): requests = [ rdf_objects.MessageHandlerRequest(client_id="C.1000000000000000", handler_name="Testhandler", request_id=i * 100, request=rdfvalue.RDFInteger(i)) for i in range(10) ] lease_time = rdfvalue.Duration("5m") with test_lib.FakeTime( rdfvalue.RDFDatetime.FromSecondsSinceEpoch(10000)): self.db.WriteMessageHandlerRequests(requests) leased = queue.Queue() self.db.RegisterMessageHandler(leased.put, lease_time, limit=5) got = [] while len(got) < 10: try: l = leased.get(True, timeout=6) except queue.Empty: self.fail( "Timed out waiting for messages, expected 10, got %d" % len(got)) self.assertLessEqual(len(l), 5) for m in l: self.assertEqual(m.leased_by, utils.ProcessIdString()) self.assertGreater(m.leased_until, rdfvalue.RDFDatetime.Now()) self.assertLess(m.timestamp, rdfvalue.RDFDatetime.Now()) m.leased_by = None m.leased_until = None m.timestamp = None got += l self.db.DeleteMessageHandlerRequests(got) got.sort(key=lambda req: req.request_id) self.assertEqual(requests, got)
def testMessageHandlers(self): client_id = self.SetupClient(100) done = threading.Event() def handle(l): worker_lib.ProcessMessageHandlerRequests(l) done.set() data_store.REL_DB.RegisterMessageHandler( handle, worker_lib.GRRWorker.message_handler_lease_time, limit=1000) data_store.REL_DB.WriteMessageHandlerRequests([ rdf_objects.MessageHandlerRequest( client_id=client_id, handler_name="StatsHandler", request_id=12345, request=rdf_client_stats.ClientStats(RSS_size=1234)) ]) self.assertTrue(done.wait(10)) results = data_store.REL_DB.ReadClientStats( client_id=client_id, min_timestamp=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(0), max_timestamp=rdfvalue.RDFDatetime.Now()) self.assertLen(results, 1) stats = results[0] self.assertEqual(stats.RSS_size, 1234) data_store.REL_DB.UnregisterMessageHandler(timeout=60) # Make sure there are no leftover requests. self.assertEqual(data_store.REL_DB.ReadMessageHandlerRequests(), [])
def ReceiveMessages(self, client_id: str, messages: Iterable[rdf_flows.GrrMessage]): """Receives and processes the messages. 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() unprocessed_msgs = [] worker_message_handler_requests = [] frontend_message_handler_requests = [] dropped_count = 0 msgs_by_session_id = collection.Group(messages, lambda m: m.session_id) for session_id, msgs in msgs_by_session_id.items(): for msg in msgs: if (msg.auth_state != msg.AuthorizationState.AUTHENTICATED and msg.session_id != self.unauth_allowed_session_id): dropped_count += 1 continue session_id_str = str(session_id) if session_id_str in message_handlers.session_id_map: request = rdf_objects.MessageHandlerRequest( client_id=msg.source.Basename(), handler_name=message_handlers. session_id_map[session_id], request_id=msg.response_id or random.UInt32(), request=msg.payload) if request.handler_name in self._SHORTCUT_HANDLERS: frontend_message_handler_requests.append(request) else: worker_message_handler_requests.append(request) elif session_id_str in self.legacy_well_known_session_ids: logging.debug( "Dropping message for legacy well known session id %s", session_id) else: unprocessed_msgs.append(msg) if dropped_count: logging.info("Dropped %d unauthenticated messages for %s", dropped_count, client_id) if unprocessed_msgs: flow_responses = [] for message in unprocessed_msgs: try: flow_responses.append( rdf_flow_objects.FlowResponseForLegacyResponse( message)) except ValueError as e: logging.warning( "Failed to parse legacy FlowResponse:\n%s\n%s", e, message) data_store.REL_DB.WriteFlowResponses(flow_responses) for msg in unprocessed_msgs: if msg.type == rdf_flows.GrrMessage.Type.STATUS: 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=msg.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) if worker_message_handler_requests: data_store.REL_DB.WriteMessageHandlerRequests( worker_message_handler_requests) if frontend_message_handler_requests: worker_lib.ProcessMessageHandlerRequests( frontend_message_handler_requests) logging.debug("Received %s messages from %s in %s sec", len(messages), client_id, time.time() - now)
def ReceiveMessagesRelationalFlows(self, client_id, messages): """Receives and processes messages for flows stored in the relational db. Args: client_id: The client which sent the messages. messages: A list of GrrMessage RDFValues. """ now = time.time() unprocessed_msgs = [] message_handler_requests = [] dropped_count = 0 for session_id, msgs in iteritems( collection.Group(messages, operator.attrgetter("session_id"))): # Remove and handle messages to WellKnownFlows leftover_msgs = self.HandleWellKnownFlows(msgs) for msg in leftover_msgs: if (msg.auth_state != msg.AuthorizationState.AUTHENTICATED and msg.session_id != self.unauth_allowed_session_id): dropped_count += 1 continue if session_id in queue_manager.session_id_map: message_handler_requests.append( rdf_objects.MessageHandlerRequest( client_id=msg.source.Basename(), handler_name=queue_manager. session_id_map[session_id], request_id=msg.response_id, request=msg.payload)) elif session_id in self.legacy_well_known_session_ids: logging.debug( "Dropping message for legacy well known session id %s", session_id) else: unprocessed_msgs.append(msg) if dropped_count: logging.info("Dropped %d unauthenticated messages for %s", dropped_count, client_id) if unprocessed_msgs: flow_responses = [] for message in unprocessed_msgs: flow_responses.append( rdf_flow_objects.FlowResponseForLegacyResponse(message)) data_store.REL_DB.WriteFlowResponses(flow_responses) for msg in unprocessed_msgs: if msg.type == rdf_flows.GrrMessage.Type.STATUS: 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=msg.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) if message_handler_requests: data_store.REL_DB.WriteMessageHandlerRequests( message_handler_requests) logging.debug("Received %s messages from %s in %s sec", len(messages), client_id, time.time() - now)
def Flush(self): """Writes the changes in this object to the datastore.""" if data_store.RelationalDBReadEnabled(category="message_handlers"): message_handler_requests = [] leftover_responses = [] for r, timestamp in self.response_queue: if r.request_id == 0 and r.session_id in session_id_map: message_handler_requests.append( rdf_objects.MessageHandlerRequest( client_id=r.source and r.source.Basename(), handler_name=session_id_map[r.session_id], request_id=r.response_id, request=r.payload)) else: leftover_responses.append((r, timestamp)) if message_handler_requests: data_store.REL_DB.WriteMessageHandlerRequests( message_handler_requests) self.response_queue = leftover_responses self.data_store.StoreRequestsAndResponses( new_requests=self.request_queue, new_responses=self.response_queue, requests_to_delete=self.requests_to_delete) # We need to make sure that notifications are written after the requests so # we flush after writing all requests and only notify afterwards. mutation_pool = self.data_store.GetMutationPool() with mutation_pool: if data_store.RelationalDBReadEnabled(category="client_messages"): if self.client_messages_to_delete: data_store.REL_DB.DeleteClientMessages( list(itervalues(self.client_messages_to_delete))) else: messages_by_queue = utils.GroupBy( list(itervalues(self.client_messages_to_delete)), lambda request: request.queue) for queue, messages in iteritems(messages_by_queue): self.Delete(queue, messages, mutation_pool=mutation_pool) if self.new_client_messages: for timestamp, messages in iteritems( utils.GroupBy(self.new_client_messages, lambda x: x[1])): self.Schedule([x[0] for x in messages], timestamp=timestamp, mutation_pool=mutation_pool) if self.notifications: for notification in itervalues(self.notifications): self.NotifyQueue(notification, mutation_pool=mutation_pool) mutation_pool.Flush() self.request_queue = [] self.response_queue = [] self.requests_to_delete = [] self.client_messages_to_delete = {} self.notifications = {} self.new_client_messages = []
def testEnrollment(self): """Test the http response to unknown clients.""" self._ClearClient() # Now communicate with the server. self.SendToServer() status = self.client_communicator.RunOnce() # We expect to receive a 406 and all client messages will be tagged as # UNAUTHENTICATED. self.assertEqual(status.code, 406) self.assertLen(self.messages, 10) self.assertEqual( self.messages[0].auth_state, rdf_flows.GrrMessage.AuthorizationState.UNAUTHENTICATED) # The next request should be an enrolling request. status = self.client_communicator.RunOnce() self.assertLen(self.messages, 11) enrolment_messages = [] for m in self.messages: if m.session_id == ca_enroller.Enroler.well_known_session_id: enrolment_messages.append(m) self.assertLen(enrolment_messages, 1) # Now we manually run the enroll well known flow with the enrollment # request. This will start a new flow for enrolling the client, sign the # cert and add it to the data store. if data_store.AFF4Enabled(): flow_obj = ca_enroller.Enroler( ca_enroller.Enroler.well_known_session_id, mode="rw", token=self.token) flow_obj.ProcessMessage(enrolment_messages[0]) else: handler = ca_enroller.EnrolmentHandler() req = rdf_objects.MessageHandlerRequest( client_id=self.client_id, request=enrolment_messages[0].payload) handler.ProcessMessages([req]) # The next client communication should be enrolled now. status = self.client_communicator.RunOnce() self.assertEqual(status.code, 200) # There should be a cert for the client right now. if data_store.AFF4Enabled(): client = aff4.FACTORY.Create(self.client_cn, aff4_grr.VFSGRRClient, mode="rw", token=self.token) self.assertTrue(client.Get(client.Schema.CERT)) else: md = data_store.REL_DB.ReadClientMetadata(self.client_id) self.assertTrue(md.certificate) # Now communicate with the server once again. self.SendToServer() status = self.client_communicator.RunOnce() self.assertEqual(status.code, 200)