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