test/TinyHTTPProxy.py

   1 #!/bin/sh -
   2 "exec" "python" "-O" "$0" "$@"
   3
   4 __doc__ = """Tiny HTTP Proxy.
   5
   6 This module implements GET, HEAD, POST, PUT and DELETE methods
   7 on BaseHTTPServer, and behaves as an HTTP proxy.  The CONNECT
   8 method is also implemented experimentally, but has not been
   9 tested yet.
  10
  11 Any help will be greatly appreciated.           SUZUKI Hisao
  12 """
  13
  14 __version__ = "0.2.1"
  15
  16 import BaseHTTPServer, select, socket, SocketServer, urlparse
  17
  18 class ProxyHandler (BaseHTTPServer.BaseHTTPRequestHandler):
  19     __base = BaseHTTPServer.BaseHTTPRequestHandler
  20     __base_handle = __base.handle
  21
  22     server_version = "TinyHTTPProxy/" + __version__
  23     rbufsize = 0                        # self.rfile Be unbuffered
  24
  25     def handle(self):
  26         (ip, port) =  self.client_address
  27         if hasattr(self, 'allowed_clients') and ip not in self.allowed_clients:
  28             self.raw_requestline = self.rfile.readline()
  29             if self.parse_request(): self.send_error(403)
  30         else:
  31             self.__base_handle()
  32
  33     def _connect_to(self, netloc, soc):
  34         i = netloc.find(':')
  35         if i >= 0:
  36             host_port = netloc[:i], int(netloc[i+1:])
  37         else:
  38             host_port = netloc, 80
  39         print "\t" "connect to %s:%d" % host_port
  40         try: soc.connect(host_port)
  41         except socket.error, arg:
  42             try: msg = arg[1]
  43             except: msg = arg
  44             self.send_error(404, msg)
  45             return 0
  46         return 1
  47
  48     def do_CONNECT(self):
  49         soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  50         try:
  51             if self._connect_to(self.path, soc):
  52                 self.log_request(200)
  53                 self.wfile.write(self.protocol_version +
  54                                  " 200 Connection established\r\n")
  55                 self.wfile.write("Proxy-agent: %s\r\n" % self.version_string())
  56                 self.wfile.write("\r\n")
  57                 self._read_write(soc, 300)
  58         finally:
  59             print "\t" "bye"
  60             soc.close()
  61             self.connection.close()
  62
  63     def do_GET(self):
  64         (scm, netloc, path, params, query, fragment) = urlparse.urlparse(
  65             self.path, 'http')
  66         if scm != 'http' or fragment or not netloc:
  67             self.send_error(400, "bad url %s" % self.path)
  68             return
  69         soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  70         try:
  71             if self._connect_to(netloc, soc):
  72                 self.log_request()
  73                 soc.send("%s %s %s\r\n" % (
  74                     self.command,
  75                     urlparse.urlunparse(('', '', path, params, query, '')),
  76                     self.request_version))
  77                 self.headers['Connection'] = 'close'
  78                 del self.headers['Proxy-Connection']
  79                 for key_val in self.headers.items():
  80                     soc.send("%s: %s\r\n" % key_val)
  81                 soc.send("\r\n")
  82                 self._read_write(soc)
  83         finally:
  84             print "\t" "bye"
  85             soc.close()
  86             self.connection.close()
  87
  88     def _read_write(self, soc, max_idling=20):
  89         iw = [self.connection, soc]
  90         ow = []
  91         count = 0
  92         while 1:
  93             count += 1
  94             (ins, _, exs) = select.select(iw, ow, iw, 3)
  95             if exs: break
  96             if ins:
  97                 for i in ins:
  98                     if i is soc:
  99                         out = self.connection
 100                     else:
 101                         out = soc
 102                     data = i.recv(8192)
 103                     if data:
 104                         out.send(data)
 105                         count = 0
 106             else:
 107                 print "\t" "idle", count
 108             if count == max_idling: break
 109
 110     do_HEAD = do_GET
 111     do_POST = do_GET
 112     do_PUT  = do_GET
 113     do_DELETE=do_GET
 114
 115 class ThreadingHTTPServer (SocketServer.ThreadingMixIn,
 116                            BaseHTTPServer.HTTPServer): pass
 117
 118 if __name__ == '__main__':
 119     from sys import argv
 120     if argv[1:] and argv[1] in ('-h', '--help'):
 121         print argv[0], "[port [allowed_client_name ...]]"
 122     else:
 123         if argv[2:]:
 124             allowed = []
 125             for name in argv[2:]:
 126                 client = socket.gethostbyname(name)
 127                 allowed.append(client)
 128                 print "Accept: %s (%s)" % (client, name)
 129             ProxyHandler.allowed_clients = allowed
 130             del argv[2:]
 131         else:
 132             print "Any clients will be served..."
 133         BaseHTTPServer.test(ProxyHandler, ThreadingHTTPServer)