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