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