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