Add a '--host' option to testserver.py .

net/tools/testserver/testserver.py

 #!/usr/bin/env python
 # Copyright (c) 2012 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.
 
 """This is a simple HTTP/FTP/SYNC/TCP/UDP/ server used for testing Chrome.
 
 It supports several test URLs, as specified by the handlers in TestPageHandler.
 By default, it listens on an ephemeral port and sends the port number back to
 the originating process over a pipe. The originating process can specify an
@@ skipping 64 matching lines @@
   def __init__(self):
     self.log = []
 
   def __getitem__(self, sessionID):
     self.log.append(('lookup', sessionID))
     raise KeyError()
 
   def __setitem__(self, sessionID, session):
     self.log.append(('insert', sessionID))
 
+
+class ClientRestrictingServerMixIn:
+  """Implements verify_request to limit connections to our configured IP
+  address."""
+
+  def verify_request(self, request, client_address):
+    return client_address[0] == self.server_address[0]
+
+
 class StoppableHTTPServer(BaseHTTPServer.HTTPServer):
-  """This is a specialization of of BaseHTTPServer to allow it
+  """This is a specialization of BaseHTTPServer to allow it
   to be exited cleanly (by setting its "stop" member to True)."""
 
   def serve_forever(self):
     self.stop = False
     self.nonce_time = None
     while not self.stop:
       self.handle_request()
     self.socket.close()
 
-class HTTPSServer(tlslite.api.TLSSocketServerMixIn, StoppableHTTPServer):
-  """This is a specialization of StoppableHTTPerver that add https support."""
+
+class HTTPServer(ClientRestrictingServerMixIn, StoppableHTTPServer):
+  """This is a specialization of StoppableHTTPerver that add client

[Paweł Hajdan Jr., 2012/02/15 17:11:01]

nit: "that adds", note the "s".

+  verification."""
+
+  pass
+
+
+class HTTPSServer(tlslite.api.TLSSocketServerMixIn,
+                  ClientRestrictingServerMixIn,
+                  StoppableHTTPServer):
+  """This is a specialization of StoppableHTTPerver that add https support and
+  client verification."""
 
   def __init__(self, server_address, request_hander_class, cert_path,
                ssl_client_auth, ssl_client_cas, ssl_bulk_ciphers,
                record_resume_info):
     s = open(cert_path).read()
     x509 = tlslite.api.X509()
     x509.parse(s)
     self.cert_chain = tlslite.api.X509CertChain([x509])
     s = open(cert_path).read()
     self.private_key = tlslite.api.parsePEMKey(s, private=True)
@@ skipping 28 matching lines @@
       tlsConnection.ignoreAbruptClose = True
       return True
     except tlslite.api.TLSAbruptCloseError:
       # Ignore abrupt close.
       return True
     except tlslite.api.TLSError, error:
       print "Handshake failure:", str(error)
       return False
 
 
-class SyncHTTPServer(StoppableHTTPServer):
+class SyncHTTPServer(ClientRestrictingServerMixIn, StoppableHTTPServer):
   """An HTTP server that handles sync commands."""
 
   def __init__(self, server_address, request_handler_class):
     # We import here to avoid pulling in chromiumsync's dependencies
     # unless strictly necessary.
     import chromiumsync
     import xmppserver
     StoppableHTTPServer.__init__(self, server_address, request_handler_class)
     self._sync_handler = chromiumsync.TestServer()
     self._xmpp_socket_map = {}
@@ skipping 84 matching lines @@
 
       for fd in write_fds:
         HandleXmppSocket(fd, self._xmpp_socket_map,
                          asyncore.dispatcher.handle_write_event)
 
       for fd in exceptional_fds:
         HandleXmppSocket(fd, self._xmpp_socket_map,
                          asyncore.dispatcher.handle_expt_event)
 
 
-class TCPEchoServer(SocketServer.TCPServer):
+class FTPServer(ClientRestrictingServerMixIn, pyftpdlib.ftpserver.FTPServer):
+  """This is a specialization of FTPServer that add client

[Paweł Hajdan Jr., 2012/02/15 17:11:01]

nit: "that adds"

+  verification."""
+
+  pass
+
+
+class TCPEchoServer(ClientRestrictingServerMixIn, SocketServer.TCPServer):
   """A TCP echo server that echoes back what it has received."""
 
   def server_bind(self):
     """Override server_bind to store the server name."""
     SocketServer.TCPServer.server_bind(self)
     host, port = self.socket.getsockname()[:2]
     self.server_name = socket.getfqdn(host)
     self.server_port = port
 
   def serve_forever(self):
     self.stop = False
     self.nonce_time = None
     while not self.stop:
       self.handle_request()
     self.socket.close()
 
 
-class UDPEchoServer(SocketServer.UDPServer):
+class UDPEchoServer(ClientRestrictingServerMixIn, SocketServer.UDPServer):
   """A UDP echo server that echoes back what it has received."""
 
   def server_bind(self):
     """Override server_bind to store the server name."""
     SocketServer.UDPServer.server_bind(self)
     host, port = self.socket.getsockname()[:2]
     self.server_name = socket.getfqdn(host)
     self.server_port = port
 
   def serve_forever(self):
@@ skipping 1301 matching lines @@
     test_name = "/chromiumsync/command"
     if not self._ShouldHandleRequest(test_name):
       return False
 
     length = int(self.headers.getheader('content-length'))
     raw_request = self.rfile.read(length)
     http_response = 200
     raw_reply = None
     if not self.server.GetAuthenticated():
       http_response = 401
-      challenge = 'GoogleLogin realm="http://127.0.0.1", service="chromiumsync"'
+      challenge = 'GoogleLogin realm="http://%s", service="chromiumsync"' % (
+        self.server.server_address[0])
     else:
       http_response, raw_reply = self.server.HandleCommand(
           self.path, raw_request)
 
     ### Now send the response to the client. ###
     self.send_response(http_response)
     if http_response == 401:
       self.send_header('www-Authenticate', challenge)
     self.end_headers()
     self.wfile.write(raw_reply)
@@ skipping 226 matching lines @@
 
 def main(options, args):
   logfile = open('testserver.log', 'w')
   sys.stderr = FileMultiplexer(sys.stderr, logfile)
   if options.log_to_console:
     sys.stdout = FileMultiplexer(sys.stdout, logfile)
   else:
     sys.stdout = logfile
 
   port = options.port
+  host = options.host
 
   server_data = {}
+  server_data['host'] = host
 
   if options.server_type == SERVER_HTTP:
     if options.cert:
       # let's make sure the cert file exists.
       if not os.path.isfile(options.cert):
         print 'specified server cert file not found: ' + options.cert + \
               ' exiting...'
         return
       for ca_cert in options.ssl_client_ca:
         if not os.path.isfile(ca_cert):
           print 'specified trusted client CA file not found: ' + ca_cert + \
                 ' exiting...'
           return
-      server = HTTPSServer(('127.0.0.1', port), TestPageHandler, options.cert,
+      server = HTTPSServer((host, port), TestPageHandler, options.cert,
                            options.ssl_client_auth, options.ssl_client_ca,
                            options.ssl_bulk_cipher, options.record_resume)
-      print 'HTTPS server started on port %d...' % server.server_port
+      print 'HTTPS server started on %s:%d...' % (host, server.server_port)
     else:
-      server = StoppableHTTPServer(('127.0.0.1', port), TestPageHandler)
-      print 'HTTP server started on port %d...' % server.server_port
+      server = HTTPServer((host, port), TestPageHandler)
+      print 'HTTP server started on %s:%d...' % (host, server.server_port)
 
     server.data_dir = MakeDataDir()
     server.file_root_url = options.file_root_url
     server_data['port'] = server.server_port
     server._device_management_handler = None
     server.policy_keys = options.policy_keys
     server.policy_user = options.policy_user
   elif options.server_type == SERVER_SYNC:
-    server = SyncHTTPServer(('127.0.0.1', port), SyncPageHandler)
+    server = SyncHTTPServer((host, port), SyncPageHandler)
     print 'Sync HTTP server started on port %d...' % server.server_port
     print 'Sync XMPP server started on port %d...' % server.xmpp_port
     server_data['port'] = server.server_port
     server_data['xmpp_port'] = server.xmpp_port
   elif options.server_type == SERVER_TCP_ECHO:
     # Used for generating the key (randomly) that encodes the "echo request"
     # message.
     random.seed()
-    server = TCPEchoServer(('127.0.0.1', port), TCPEchoHandler)
+    server = TCPEchoServer((host, port), TCPEchoHandler)
     print 'Echo TCP server started on port %d...' % server.server_port
     server_data['port'] = server.server_port
   elif options.server_type == SERVER_UDP_ECHO:
     # Used for generating the key (randomly) that encodes the "echo request"
     # message.
     random.seed()
-    server = UDPEchoServer(('127.0.0.1', port), UDPEchoHandler)
+    server = UDPEchoServer((host, port), UDPEchoHandler)
     print 'Echo UDP server started on port %d...' % server.server_port
     server_data['port'] = server.server_port
   # means FTP Server
   else:
     my_data_dir = MakeDataDir()
 
     # Instantiate a dummy authorizer for managing 'virtual' users
     authorizer = pyftpdlib.ftpserver.DummyAuthorizer()
 
     # Define a new user having full r/w permissions and a read-only
     # anonymous user
     authorizer.add_user('chrome', 'chrome', my_data_dir, perm='elradfmw')
 
     authorizer.add_anonymous(my_data_dir)
 
     # Instantiate FTP handler class
     ftp_handler = pyftpdlib.ftpserver.FTPHandler
     ftp_handler.authorizer = authorizer
 
     # Define a customized banner (string returned when client connects)
     ftp_handler.banner = ("pyftpdlib %s based ftpd ready." %
                           pyftpdlib.ftpserver.__ver__)
 
-    # Instantiate FTP server class and listen to 127.0.0.1:port
-    address = ('127.0.0.1', port)
-    server = pyftpdlib.ftpserver.FTPServer(address, ftp_handler)
+    # Instantiate FTP server class and listen to address:port
+    server = pyftpdlib.ftpserver.FTPServer((host, port), ftp_handler)
     server_data['port'] = server.socket.getsockname()[1]
     print 'FTP server started on port %d...' % server_data['port']
 
   # Notify the parent that we've started. (BaseServer subclasses
   # bind their sockets on construction.)
   if options.startup_pipe is not None:
     server_data_json = json.dumps(server_data)
     server_data_len = len(server_data_json)
     print 'sending server_data: %s (%d bytes)' % (
       server_data_json, server_data_len)
@@ skipping 81 matching lines @@
                            'the server. If ther server has multiple keys, it '
                            'will rotate through them in at each request a '
                            'round-robin fashion. The server will generate a '
                            'random key if none is specified on the command '
                            'line.')
   option_parser.add_option('', '--policy-user', default='user@example.com',
                            dest='policy_user',
                            help='Specify the user name the server should '
                            'report back to the client as the user owning the '
                            'token used for making the policy request.')
+  option_parser.add_option('', '--host', default='127.0.0.1',
+                           dest='host',
+                           help='Hostname or IP upon which the server will '
+                           'listen. Client connections will also only be '
+                           'allowed from this address.')
   options, args = option_parser.parse_args()
 
   sys.exit(main(options, args))