def ProcessMessageHandlerRequests(requests): """Processes message handler requests.""" logging.debug("Leased message handler request ids: %s", ",".join(str(r.request_id) for r in requests)) grouped_requests = collection.Group(requests, lambda r: r.handler_name) for handler_name, requests_for_handler in iteritems(grouped_requests): handler_cls = handler_registry.handler_name_map.get(handler_name) if not handler_cls: logging.error("Unknown message handler: %s", handler_name) continue num_requests = len(requests_for_handler) WELL_KNOWN_FLOW_REQUESTS.Increment(fields=[handler_name], delta=num_requests) try: logging.debug("Running %d messages for handler %s", num_requests, handler_name) handler_cls().ProcessMessages(requests_for_handler) except Exception as e: # pylint: disable=broad-except logging.exception( "Exception while processing message handler %s: %s", handler_name, e) logging.debug("Deleting message handler request ids: %s", ",".join(str(r.request_id) for r in requests)) data_store.REL_DB.DeleteMessageHandlerRequests(requests)
def _DownloadFiles(self, responses): if not responses.success: self.Log("Failed to run ArtifactCollectorFlow: %s", responses.status) return results_with_pathspecs = [] results_without_pathspecs = [] for response in responses: pathspecs = self._FindMatchingPathspecs(response) if pathspecs: for pathspec in pathspecs: result = ArtifactFilesDownloaderResult( original_result_type=response.__class__.__name__, original_result=response, found_pathspec=pathspec) results_with_pathspecs.append(result) else: result = ArtifactFilesDownloaderResult( original_result_type=response.__class__.__name__, original_result=response) results_without_pathspecs.append(result) grouped_results = collection.Group( results_with_pathspecs, lambda x: x.found_pathspec.CollapsePath()) for _, group in grouped_results.items(): self.StartFileFetch( group[0].found_pathspec, request_data=dict(results=group)) for result in results_without_pathspecs: self.SendReply(result)
def _ProcessMessageHandlerRequests(self, requests): """Processes message handler requests.""" logging.debug("Leased message handler request ids: %s", ",".join(str(r.request_id) for r in requests)) grouped_requests = collection.Group(requests, lambda r: r.handler_name) for handler_name, requests_for_handler in iteritems(grouped_requests): handler_cls = handler_registry.handler_name_map.get(handler_name) if not handler_cls: logging.error("Unknown message handler: %s", handler_name) continue stats_collector_instance.Get().IncrementCounter( "well_known_flow_requests", fields=[handler_name]) try: logging.debug("Running %d messages for handler %s", len(requests_for_handler), handler_name) handler_cls(token=self.token).ProcessMessages(requests_for_handler) except Exception as e: # pylint: disable=broad-except logging.exception("Exception while processing message handler %s: %s", handler_name, e) logging.debug("Deleting message handler request ids: %s", ",".join(str(r.request_id) for r in requests)) data_store.REL_DB.DeleteMessageHandlerRequests(requests)
def Flush(self): """Writes the changes in this object to the datastore.""" 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: messages_by_queue = collection.Group( 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( collection.Group(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 testGenerator(self): def Generate(): yield 4 yield 8 yield 15 yield 16 yield 23 yield 42 result = collection.Group(Generate(), key=lambda num: num % 2) expected = {0: [4, 8, 16, 42], 1: [15, 23]} self.assertEqual(result, expected)
def MultiNotifyQueue(self, notifications, mutation_pool=None): """This is the same as NotifyQueue but for several session_ids at once. Args: notifications: A list of notifications. mutation_pool: A MutationPool object to schedule Notifications on. Raises: RuntimeError: An invalid session_id was passed. """ extract_queue = lambda notification: notification.session_id.Queue() for queue, notifications in iteritems( collection.Group(notifications, extract_queue)): self._MultiNotifyQueue(queue, notifications, mutation_pool=mutation_pool)
def Schedule(self, tasks, mutation_pool, timestamp=None): """Schedule a set of Task() instances.""" non_fleetspeak_tasks = [] for queue, queued_tasks in iteritems( collection.Group(tasks, lambda x: x.queue)): if not queue: continue client_id = _GetClientIdFromQueue(queue) if fleetspeak_utils.IsFleetspeakEnabledClient(client_id, token=self.token): for task in queued_tasks: fleetspeak_utils.SendGrrMessageThroughFleetspeak( client_id, task) continue non_fleetspeak_tasks.extend(queued_tasks) timestamp = timestamp or self.frozen_timestamp mutation_pool.QueueScheduleTasks(non_fleetspeak_tasks, timestamp)
def ConvertValuesWithMetadata(metadata_value_pairs, options=None): """Converts a set of RDFValues into a set of export-friendly RDFValues. Args: metadata_value_pairs: Tuples of (metadata, rdf_value), where metadata is an instance of ExportedMetadata and rdf_value is an RDFValue subclass instance to be exported. options: rdfvalue.ExportOptions instance that will be passed to ExportConverters. Yields: Converted values. Converted values may be of different types. Raises: NoConverterFound: in case no suitable converters were found for a value in metadata_value_pairs. This error is only raised after all values in metadata_value_pairs are attempted to be converted. If there are multiple value types that could not be converted because of the lack of corresponding converters, only the last one will be specified in the exception message. """ no_converter_found_error = None metadata_value_groups = collection.Group( metadata_value_pairs, lambda pair: pair[1].__class__.__name__) for metadata_values_group in metadata_value_groups.values(): _, first_value = metadata_values_group[0] converters_classes = export_converters_registry.GetConvertersByValue( first_value) if not converters_classes: no_converter_found_error = "No converters found for value: %s" % str( first_value) continue converters = [cls(options) for cls in converters_classes] for converter in converters: for result in converter.BatchConvert(metadata_values_group): yield result if no_converter_found_error is not None: raise NoConverterFound(no_converter_found_error)
def DeleteNotifications(self, session_ids, start=None, end=None): """This deletes the notification when all messages have been processed.""" if not session_ids: return for session_id in session_ids: if not isinstance(session_id, rdfvalue.SessionID): raise RuntimeError( "Can only delete notifications for rdfvalue.SessionIDs.") if start is None: start = 0 else: start = int(start) if end is None: end = self.frozen_timestamp or rdfvalue.RDFDatetime.Now() for queue, ids in iteritems( collection.Group(session_ids, lambda session_id: session_id.Queue())): queue_shards = self.GetAllNotificationShards(queue) self.data_store.DeleteNotifications(queue_shards, ids, start, end)
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 testByFirstLetter(self): result = collection.Group(["foo", "bar", "baz"], key=lambda text: text[0]) expected = {"f": ["foo"], "b": ["bar", "baz"]} self.assertEqual(result, expected)
def testByIdentity(self): result = collection.Group([3, 2, 1, 1, 5, 3, 1, 5], key=lambda num: num) expected = {1: [1, 1, 1], 2: [2], 3: [3, 3], 5: [5, 5]} self.assertEqual(result, expected)
def testEmpty(self): result = collection.Group([], key=lambda _: None) expected = {} self.assertEqual(result, expected)
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. """ if data_store.RelationalDBEnabled(): return self.ReceiveMessagesRelationalFlows(client_id, messages) now = time.time() with queue_manager.QueueManager(token=self.token) as manager: for session_id, msgs in iteritems( collection.Group(messages, operator.attrgetter("session_id"))): # 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) # 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, 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 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 = collection.Group( 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( collection.Group(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 = []