Adding the initial code for Omnibot multi-machine support

testing/legion/discovery_server.py

+# Copyright 2015 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+"""The discovery server used to register clients.
+
+The discovery server is started by the host controller and allows the clients
+to register themselves when they start. Authentication of the client controllers
+is based on an OTP passed to the client controller binary on startup.
+"""
+import logging
+import threading
+import xmlrpclib
+import SimpleXMLRPCServer
+
+SERVER_ADDRESS = ''
+SERVER_PORT = 31710
+
+
+class DiscoveryServer(object):
+  """Discovery server run on the host."""
+
+  def __init__(self):
+    self._expected_clients = {}
+    self._rpc_server = None
+    self._thread = None
+
+  def _RegisterClientRPC(self, otp, ip):
+    """The RPC used by a client to register with the discovery server.
+
+    Args:
+      otp: The one time token issued by the host.
+      ip: The ip address of the client.
+
+    Raises:
+      KeyError if the OTP isn't found in _expected_clients
+    """
+    if otp not in self._expected_clients:
+      raise KeyError('OTP not found')
+    cb = self._expected_clients[otp]
+    del self._expected_clients[otp]
+    cb(ip)
+
+  def RegisterClientCallback(self, otp, callback):
+    """Registers a callback associated with an OTP.
+
+    Args:
+      otp: A one time token used by the client to authenticate.
+      callback: The callback used when the client connects.
+
+    Raises:
+      TypeError if the callback is not callable.
+    """
+    if not callable(callback):
+      raise TypeError('callback is not callable')
+    self._expected_clients[otp] = callback
+
+  def Start(self, address=None, port=None):
+    """Starts the discovery server.
+
+    Args:
+      address: The address to run the server on ('', 'localhost', etc).
+      port: The port to run the server on.
+    """
+    address = address or SERVER_ADDRESS
+    port = port = SERVER_PORT
+    logging.debug('Starting discovery server')
+    self._rpc_server = SimpleXMLRPCServer.SimpleXMLRPCServer(
+        (address, port), allow_none=True, logRequests=False)
+    self._rpc_server.register_function(
+        self._RegisterClientRPC, 'RegisterClient')
+    self._thread = threading.Thread(target=self._rpc_server.serve_forever)
+    self._thread.start()
+
+  def Shutdown(self):
+    """Shuts the discovery server down."""
+    if self._thread and self._thread.is_alive():
+      logging.debug('Shutting down discovery server')
+      self._rpc_server.shutdown()
+
+  @staticmethod
+  def Connect(host):
+    """Create a connection to the discovery server."""
+    server = 'http://%s:%s' % (host, SERVER_PORT)
+    logging.info('Connecting to a discovery server at %s', server)
+    return xmlrpclib.Server(server)