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         if not hasattr(self.client, "proxy"):
 131             self.client.proxy = ProxyConnection()
 132             self.client.proxy.connect()
 133
 134         multitask.add(self.proxy_relay(reqtype))
 135
 136         return True
 137
 138     def onPOST(self, client, *args):
 139         """Serve a POST request."""
 140         return self.on_query(client, "POST", *args)
 141
 142     def onGET(self, client, *args):
 143         """Serve a GET request."""
 144         return self.on_query(client, "GET", *args)
 145
 146     def proxy_relay(self, reqtype):
 147
 148         p = self.client.proxy
 149
 150         # send command
 151         req = "%s %s %s\n" % (reqtype, self.hr.path, self.hr.request_version)
 152         print "req", req
 153         yield p.ss.write(req)
 154
 155         # send headers
 156         hdrs = str(self.hr.headers)
 157         print "hdrs", hdrs
 158         yield p.ss.write(hdrs)
 159         yield p.ss.write('\r\n')
 160
 161         conntype = self.hr.headers.get('Connection', "")
 162         keepalive = conntype.lower() == 'keep-alive'
 163
 164         # now content
 165         if self.hr.headers.has_key('content-length'):
 166             max_chunk_size = 10*1024*1024
 167             size_remaining = int(self.hr.headers["content-length"])
 168             L = []
 169             print "size_remaining", size_remaining
 170             while size_remaining:
 171                 chunk_size = min(size_remaining, max_chunk_size)
 172                 data = self.hr.rfile.read(chunk_size)
 173                 print "proxy rfile read", repr(data)
 174                 yield multitask.send(p.sock, data)
 175                 size_remaining -= len(data)
 176
 177         # now read response and write back
 178         # HTTP/1.0 200 OK status line etc.
 179         line = (yield p.ss.readline())
 180         yield self.client.writeMessage(line)
 181
 182         res = ''
 183         try:
 184             while 1:
 185                 line = (yield p.ss.readline())
 186                 print "reading from proxy", repr(line)
 187                 res += line
 188                 if line in ['\n', '\r\n']:
 189                     break
 190         except StopIteration:
 191             if httpd._debug: print "proxy read stopiter"
 192             # TODO: close connection
 193         except:
 194             if httpd._debug:
 195                 print 'proxy read error', \
 196                       (traceback and traceback.print_exc() or None)
 197             # TODO: close connection
 198
 199         f = StringIO(res)
 200
 201         # Examine the headers and look for a Connection directive
 202         respheaders = mimetools.Message(f, 0)
 203         print "response headers", str(respheaders)
 204         remote = self.client.remote
 205         rcooks = httpd.process_cookies(respheaders, remote, "Set-Cookie", False)
 206         rcooks['session'] = self.hr.response_cookies['session'].value # nooo
 207         rcooks['session']['expires'] = \
 208                 self.hr.response_cookies['session']['expires']
 209         self.hr.response_cookies = rcooks
 210         print "rcooks", str(rcooks)
 211
 212         # send all but Set-Cookie headers
 213         del respheaders['Set-Cookie'] # being replaced
 214         yield self.client.writeMessage(str(respheaders))
 215
 216         # now replacement cookies
 217         for k, v in rcooks.items():
 218             val = v.output()
 219             yield self.client.writeMessage(val+"\r\n")
 220
 221         # check connection for "closed" header
 222         if keepalive:
 223             conntype = respheaders.get('Connection', "")
 224             if conntype.lower() == 'close':
 225                 self.hr.close_connection = 1
 226             elif (conntype.lower() == 'keep-alive' and
 227                   self.hr.protocol_version >= "HTTP/1.1"):
 228                 self.hr.close_connection = 0
 229
 230         # write rest of data
 231         print "writing to client body"
 232         yield self.client.writeMessage("\r\n")
 233
 234         if respheaders.has_key('content-length'):
 235             max_chunk_size = 10*1024*1024
 236             size_remaining = int(respheaders["content-length"])
 237             while size_remaining:
 238                 chunk_size = min(size_remaining, max_chunk_size)
 239                 data = (yield p.ss.read(chunk_size))
 240                 print "reading from proxy expecting", size_remaining, repr(data)
 241                 yield self.client.writeMessage(data)
 242                 size_remaining -= len(data)
 243         else:
 244             while True:
 245                 #data = p.read()
 246                 try:
 247                     data = (yield p.ss.read(1024))
 248                 except httpd.ConnectionClosed:
 249                     break
 250                 print "reading from proxy", repr(data)
 251                 if data == '':
 252                     break
 253                 yield self.client.writeMessage(data)
 254
 255
 256         if self.hr.close_connection:
 257             print 'proxy wants client to close_connection'
 258             try:
 259                 yield self.client.connectionClosed()
 260                 pass
 261             except httpd.ConnectionClosed:
 262                 print 'close_connection done'
 263                 pass
 264
 265         raise StopIteration
 266