ProxyServer.py

   1 """ HTTP Proxy Server
   2
   3 This module acts as an HTTP Proxy Server.  However it adds the
   4 multitaskhttpd auto-generated session cookie to the headers.
   5
   6 It also changes the connection type to use keep-alive, so that
   7 the connection stays open to the server, resulting in an
   8 "apparent" persistent connection.
   9
  10 Thus, a service on the receiving end of this proxy server may reliably
  11 keep persistent state, even though the browsers connecting to it
  12 may be doing HTTP 1.0 or have an unreliable internet connection.
  13
  14 """
  15
  16
  17 __version__ = "0.6"
  18
  19 __all__ = ["SimpleHTTPRequestHandler"]
  20
  21 import os
  22 import posixpath
  23 import BaseHTTPServer
  24 import urllib
  25 import urlparse
  26 import traceback
  27 import cgi
  28 import shutil
  29 import mimetools
  30 import multitask
  31 import socket
  32 try:
  33     from cStringIO import StringIO
  34 except ImportError:
  35     from StringIO import StringIO
  36
  37 import httpd
  38
  39 class NotConnected(Exception):
  40     pass
  41
  42 class ProxyConnection:
  43
  44     auto_open = 1
  45     debuglevel = 0
  46     strict = 0
  47
  48     def __init__(self):
  49         self.sock = None
  50         self.host = "127.0.0.1"
  51         self.port = 60001
  52
  53
  54     def connect(self):
  55         """Connect to the host and port specified in __init__."""
  56         msg = "getaddrinfo returns an empty list"
  57         for res in socket.getaddrinfo(self.host, self.port, 0,
  58                                       socket.SOCK_STREAM):
  59             af, socktype, proto, canonname, sa = res
  60             try:
  61                 self.sock = socket.socket(af, socktype, proto)
  62                 if self.debuglevel > 0:
  63                     print "connect: (%s, %s)" % (self.host, self.port)
  64                 self.sock.connect(sa)
  65             except socket.error, msg:
  66                 if self.debuglevel > 0:
  67                     print 'connect fail:', (self.host, self.port)
  68                 if self.sock:
  69                     self.sock.close()
  70                 self.sock = None
  71                 continue
  72             break
  73         if not self.sock:
  74             raise socket.error, msg
  75
  76         self.ss = httpd.SockStream(self.sock)
  77
  78     def close(self):
  79         """Close the connection to the HTTP server."""
  80         if self.sock:
  81             self.sock.close()   # close it manually... there may be other refs
  82             self.sock = None
  83
  84     def send(self, str):
  85         """Send `str' to the server."""
  86         if self.sock is None:
  87             if self.auto_open:
  88                 self.connect()
  89             else:
  90                 raise NotConnected()
  91
  92         # send the data to the server. if we get a broken pipe, then close
  93         # the socket. we want to reconnect when somebody tries to send again.
  94         #
  95         # NOTE: we DO propagate the error, though, because we cannot simply
  96         #       ignore the error... the caller will know if they can retry.
  97         if self.debuglevel > 0:
  98             print "send:", repr(str)
  99         try:
 100             yield multitask.send(self.sock, str)
 101         except socket.error, v:
 102             if v[0] == 32:      # Broken pipe
 103                 self.close()
 104             raise
 105
 106     def read(self, num_bytes=1024):
 107         data = (yield multitask.recv(self.sock, num_bytes))
 108
 109
 110 class ProxyServerRequestHandler(object):
 111
 112     """Simple HTTP request handler with GET and HEAD commands.
 113
 114     This serves files from the current directory and any of its
 115     subdirectories.  The MIME type for files is determined by
 116     calling the .guess_type() method.
 117
 118     The GET and HEAD requests are identical except that the HEAD
 119     request omits the actual contents of the file.
 120
 121     """
 122
 123     server_version = "SimpleHTTP/" + __version__
 124
 125     def on_query(self, client, reqtype, *args):
 126         """Serve a request."""
 127         self.client = client
 128         self.hr = args[0]
 129         print "on_query", reqtype, repr(self.hr.headers), str(self.hr.headers)
 130         session = self.client.session
 131         p = self.proxies.get(session, None)
 132         if not p:
 133             proxy = ProxyConnection()
 134             proxy.connect()
 135             self.proxies[session] = proxy
 136
 137         multitask.add(self.proxy_relay(reqtype))
 138
 139         return True
 140
 141     def onPOST(self, client, *args):
 142         """Serve a POST request."""
 143         return self.on_query(client, "POST", *args)
 144
 145     def onGET(self, client, *args):
 146         """Serve a GET request."""
 147         return self.on_query(client, "GET", *args)
 148
 149     def proxy_relay(self, reqtype):
 150
 151         session = self.client.session
 152         p = self.proxies[session]
 153
 154         # send command
 155         req = "%s %s %s\n" % (reqtype, self.hr.path, self.hr.request_version)
 156         print "req", req
 157         yield p.ss.write(req)
 158
 159         # send headers
 160         hdrs = str(self.hr.headers)
 161         print "hdrs", hdrs
 162         yield p.ss.write(hdrs)
 163         yield p.ss.write('\r\n')
 164
 165         conntype = self.hr.headers.get('Connection', "")
 166         keepalive = conntype.lower() == 'keep-alive'
 167
 168         # now content
 169         if self.hr.headers.has_key('content-length'):
 170             max_chunk_size = 10*1024*1024
 171             size_remaining = int(self.hr.headers["content-length"])
 172             L = []
 173             print "size_remaining", size_remaining
 174             while size_remaining:
 175                 chunk_size = min(size_remaining, max_chunk_size)
 176                 data = self.hr.rfile.read(chunk_size)
 177                 print "proxy rfile read", repr(data)
 178                 yield multitask.send(p.sock, data)
 179                 size_remaining -= len(data)
 180
 181         # now read response and write back
 182         # HTTP/1.0 200 OK status line etc.
 183         line = (yield p.ss.readline())
 184         yield self.client.writeMessage(line)
 185
 186         res = ''
 187         try:
 188             while 1:
 189                 line = (yield p.ss.readline())
 190                 print "reading from proxy", repr(line)
 191                 res += line
 192                 if line in ['\n', '\r\n']:
 193                     break
 194         except StopIteration:
 195             if httpd._debug: print "proxy read stopiter"
 196             # TODO: close connection
 197         except:
 198             if httpd._debug:
 199                 print 'proxy read error', \
 200                       (traceback and traceback.print_exc() or None)
 201             # TODO: close connection
 202
 203         f = StringIO(res)
 204
 205         # Examine the headers and look for a Connection directive
 206         respheaders = mimetools.Message(f, 0)
 207         print "response headers", str(respheaders)
 208         remote = self.client.remote
 209         rcooks = httpd.process_cookies(respheaders, remote, "Set-Cookie", False)
 210         rcooks['session'] = self.hr.response_cookies['session'].value # nooo
 211         rcooks['session']['expires'] = \
 212                 self.hr.response_cookies['session']['expires']
 213         self.hr.response_cookies = rcooks
 214         print "rcooks", str(rcooks)
 215
 216         # send all but Set-Cookie headers
 217         del respheaders['Set-Cookie'] # being replaced
 218         yield self.client.writeMessage(str(respheaders))
 219
 220         # now replacement cookies
 221         for k, v in rcooks.items():
 222             val = v.output()
 223             yield self.client.writeMessage(val+"\r\n")
 224
 225         # check connection for "closed" header
 226         if keepalive:
 227             conntype = respheaders.get('Connection', "")
 228             if conntype.lower() == 'close':
 229                 self.hr.close_connection = 1
 230             elif (conntype.lower() == 'keep-alive' and
 231                   self.hr.protocol_version >= "HTTP/1.1"):
 232                 self.hr.close_connection = 0
 233
 234         # write rest of data
 235         print "writing to client body"
 236         yield self.client.writeMessage("\r\n")
 237
 238         if respheaders.has_key('content-length'):
 239             max_chunk_size = 10*1024*1024
 240             size_remaining = int(respheaders["content-length"])
 241             while size_remaining:
 242                 chunk_size = min(size_remaining, max_chunk_size)
 243                 data = (yield p.ss.read(chunk_size))
 244                 print "reading from proxy expecting", size_remaining, repr(data)
 245                 yield self.client.writeMessage(data)
 246                 size_remaining -= len(data)
 247         else:
 248             while True:
 249                 #data = p.read()
 250                 try:
 251                     data = (yield p.ss.read(1024))
 252                 except httpd.ConnectionClosed:
 253                     break
 254                 print "reading from proxy", repr(data)
 255                 if data == '':
 256                     break
 257                 yield self.client.writeMessage(data)
 258
 259         raise StopIteration
 260
 261         if self.hr.close_connection:
 262             print 'proxy wants client to close_connection'
 263             try:
 264                 raise httpd.ConnectionClosed
 265                 yield self.client.connectionClosed()
 266             except httpd.ConnectionClosed:
 267                 print 'close_connection done'
 268                 pass
 269
 270         raise StopIteration
 271