How to use fleetspeak - 9 common examples

def Init(service_client=None):
  """Initializes the Fleetspeak connector."""
  global CONN
  global label_map

  if service_client is None:
    service_client_cls = fs_client.InsecureGRPCServiceClient

    fleetspeak_message_listen_address = (
        config.CONFIG["Server.fleetspeak_message_listen_address"] or None)
    fleetspeak_server = config.CONFIG["Server.fleetspeak_server"] or None

    if fleetspeak_message_listen_address is None and fleetspeak_server is None:
      logging.warning(
          "Missing config options `Server.fleetspeak_message_listen_address', "
          "`Server.fleetspeak_server', at least one of which is required to "
          "initialize a connection to Fleetspeak; Not using Fleetspeak.")
      return

    service_client = service_client_cls(
        "GRR",
        fleetspeak_message_listen_address=fleetspeak_message_listen_address,
        fleetspeak_server=fleetspeak_server,

def UpdateClientsFromFleetspeak(clients):
  """Updates ApiClient records to include info from Fleetspeak."""
  if not fleetspeak_connector.CONN or not fleetspeak_connector.CONN.outgoing:
    # FS not configured, or an outgoing connection is otherwise unavailable.
    return
  id_map = {}
  for client in clients:
    if client.fleetspeak_enabled:
      id_map[fleetspeak_utils.GRRIDToFleetspeakID(client.client_id)] = client
  if not id_map:
    return
  res = fleetspeak_connector.CONN.outgoing.ListClients(
      admin_pb2.ListClientsRequest(client_ids=list(iterkeys(id_map))))
  for read in res.clients:
    api_client = id_map[read.client_id]
    api_client.last_seen_at = fleetspeak_utils.TSToRDFDatetime(
        read.last_contact_time)
    api_client.last_clock = fleetspeak_utils.TSToRDFDatetime(read.last_clock)

def _GetAddrFromFleetspeak(client_id):
  res = fleetspeak_connector.CONN.outgoing.ListClients(
      admin_pb2.ListClientsRequest(
          client_ids=[fleetspeak_utils.GRRIDToFleetspeakID(client_id)]))
  if not res.clients or not res.clients[0].last_contact_address:
    return "", None
  # last_contact_address typically includes a port
  parsed = urlparse.urlparse("//{}".format(res.clients[0].last_contact_address))
  ip_str = parsed.hostname
  return ip_str, ipaddress.ip_address(ip_str)

def GetLabelsFromFleetspeak(client_id):
  """Returns labels for a Fleetspeak-enabled client.

  Fleetspeak-enabled clients delegate labeling to Fleetspeak, as opposed to
  using labels in the GRR config.

  Args:
    client_id: Id of the client to fetch Fleetspeak labels for.

  Returns:
    A list of client labels.
  """
  res = fleetspeak_connector.CONN.outgoing.ListClients(
      admin_pb2.ListClientsRequest(client_ids=[GRRIDToFleetspeakID(client_id)]))
  if not res.clients or not res.clients[0].labels:
    return []

  grr_labels = []
  label_prefix = config.CONFIG["Server.fleetspeak_label_prefix"]
  for fs_label in res.clients[0].labels:
    if (fs_label.service_name != "client" or
        (label_prefix and not fs_label.label.startswith(label_prefix))):
      continue
    try:
      grr_labels.append(fleetspeak_connector.label_map[fs_label.label])
    except KeyError:
      grr_labels.append(fs_label.label)

  return grr_labels

def _SendMessages(self, grr_msgs, background=False):
    """Sends a block of messages through Fleetspeak."""
    message_list = rdf_flows.PackedMessageList()
    communicator.Communicator.EncodeMessageList(
        rdf_flows.MessageList(job=grr_msgs), message_list)
    fs_msg = fs_common_pb2.Message(
        message_type="MessageList",
        destination=fs_common_pb2.Address(service_name="GRR"),
        background=background)
    fs_msg.data.Pack(message_list.AsPrimitiveProto())

    for grr_msg in grr_msgs:
      if (grr_msg.session_id is None or grr_msg.request_id is None or
          grr_msg.response_id is None):
        continue
      # Place all ids in a single annotation, instead of having separate
      # annotations for the flow-id, request-id and response-id. This reduces
      # overall size of the annotations by half (~60 bytes to ~30 bytes).
      annotation = fs_msg.annotations.entries.add()
      annotation.key = _DATA_IDS_ANNOTATION_KEY
      annotation.value = "%s:%d:%d" % (grr_msg.session_id.Basename(),
                                       grr_msg.request_id, grr_msg.response_id)
      if fs_msg.annotations.ByteSize() >= _MAX_ANNOTATIONS_BYTES:
        break

def SendGrrMessageThroughFleetspeak(grr_id, grr_msg):
  """Sends the given GrrMessage through FS."""
  fs_msg = fs_common_pb2.Message(
      message_type="GrrMessage",
      destination=fs_common_pb2.Address(
          client_id=GRRIDToFleetspeakID(grr_id), service_name="GRR"))
  fs_msg.data.Pack(grr_msg.AsPrimitiveProto())
  if grr_msg.session_id is not None:
    annotation = fs_msg.annotations.entries.add()
    annotation.key, annotation.value = "flow_id", grr_msg.session_id.Basename()
  if grr_msg.request_id is not None:
    annotation = fs_msg.annotations.entries.add()
    annotation.key, annotation.value = "request_id", str(grr_msg.request_id)
  fleetspeak_connector.CONN.outgoing.InsertMessage(fs_msg)

def SendGrrMessageThroughFleetspeak(grr_id, grr_msg):
  """Sends the given GrrMessage through FS."""
  fs_msg = fs_common_pb2.Message(
      message_type="GrrMessage",
      destination=fs_common_pb2.Address(
          client_id=GRRIDToFleetspeakID(grr_id), service_name="GRR"))
  fs_msg.data.Pack(grr_msg.AsPrimitiveProto())
  if grr_msg.session_id is not None:
    annotation = fs_msg.annotations.entries.add()
    annotation.key, annotation.value = "flow_id", grr_msg.session_id.Basename()
  if grr_msg.request_id is not None:
    annotation = fs_msg.annotations.entries.add()
    annotation.key, annotation.value = "request_id", str(grr_msg.request_id)
  fleetspeak_connector.CONN.outgoing.InsertMessage(fs_msg)

def _SendMessages(self, grr_msgs, background=False):
    """Sends a block of messages through Fleetspeak."""
    message_list = rdf_flows.PackedMessageList()
    communicator.Communicator.EncodeMessageList(
        rdf_flows.MessageList(job=grr_msgs), message_list)
    fs_msg = fs_common_pb2.Message(
        message_type="MessageList",
        destination=fs_common_pb2.Address(service_name="GRR"),
        background=background)
    fs_msg.data.Pack(message_list.AsPrimitiveProto())

    for grr_msg in grr_msgs:
      if (grr_msg.session_id is None or grr_msg.request_id is None or
          grr_msg.response_id is None):
        continue
      # Place all ids in a single annotation, instead of having separate
      # annotations for the flow-id, request-id and response-id. This reduces
      # overall size of the annotations by half (~60 bytes to ~30 bytes).
      annotation = fs_msg.annotations.entries.add()
      annotation.key = _DATA_IDS_ANNOTATION_KEY
      annotation.value = "%s:%d:%d" % (grr_msg.session_id.Basename(),
                                       grr_msg.request_id, grr_msg.response_id)

def __init__(self):
    self._fs = fs_client.FleetspeakConnection(
        version=config.CONFIG["Source.version_string"])

    self._sender_queue = queue.Queue(
        maxsize=GRRFleetspeakClient._SENDER_QUEUE_MAXSIZE)

    self._threads = {}

    if platform.system() == "Windows":
      internal_nanny_monitoring = False
      heart_beat_cb = self._fs.Heartbeat
    else:
      # TODO(amoser): Once the Fleetspeak nanny functionality is
      # production ready, change this to
      # internal_nanny_monitoring=False
      # heart_beat_cb=self._fs.Heartbeat
      internal_nanny_monitoring = True