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     debuglevel = 0
 125     server_version = "SimpleHTTP/" + __version__
 126
 127     def on_query(self, client, reqtype, *args):
 128         """Serve a request."""
 129
 130         self.client = client
 131         self.hr = args[0]
 132         if self.debuglevel > 0:
 133             print "on_query", reqtype, repr(self.hr.headers), \
 134                               str(self.hr.headers)
 135         session = self.client.session
 136         p = self.proxies.get(session, None)
 137         if not p:
 138             proxy = ProxyConnection()
 139             proxy.connect()
 140             self.proxies[session] = proxy
 141
 142         yield self.proxy_relay(reqtype)
 143
 144         #while p.serving:
 145         #    (yield multitask.sleep(0.01))
 146
 147         raise StopIteration
 148
 149     def onPOST(self, client, *args):
 150         """Serve a POST request."""
 151         yield self.on_query(client, "POST", *args)
 152
 153     def onGET(self, client, *args):
 154         """Serve a GET request."""
 155         yield self.on_query(client, "GET", *args)
 156
 157     def proxy_relay(self, reqtype):
 158
 159         session = self.client.session
 160         p = self.proxies[session]
 161
 162         p.serving = True
 163
 164         try:
 165             # send command
 166             req = "%s %s %s\n" % (reqtype, self.hr.path, "HTTP/1.1")
 167             if self.debuglevel > 0:
 168                 print "req", req
 169             yield p.ss.write(req)
 170
 171             conntype = self.hr.headers.get('Connection', "")
 172             keepalive = conntype.lower() == 'keep-alive'
 173
 174             self.hr.headers['Connection'] = 'keep-alive'
 175             self.hr.close_connection = 0
 176
 177             # send headers
 178             hdrs = str(self.hr.headers)
 179             if self.debuglevel > 0:
 180                 print "hdrs", hdrs
 181             yield p.ss.write(hdrs)
 182             yield p.ss.write('\r\n')
 183
 184             # now content
 185             if self.hr.headers.has_key('content-length'):
 186                 max_chunk_size = 10*1024*1024
 187                 size_remaining = int(self.hr.headers["content-length"])
 188                 L = []
 189                 if self.debuglevel > 0:
 190                     print "size_remaining", size_remaining
 191                 while size_remaining:
 192                     chunk_size = min(size_remaining, max_chunk_size)
 193                     data = self.hr.rfile.read(chunk_size)
 194                     if self.debuglevel > 0:
 195                         print "proxy rfile read", repr(data)
 196                     yield multitask.send(p.sock, data)
 197                     size_remaining -= len(data)
 198
 199             # now read response and write back
 200             # HTTP/1.0 200 OK status line etc.
 201             responseline = (yield p.ss.readline())
 202             yield self.client.writeMessage(responseline)
 203
 204             res = ''
 205             try:
 206                 while 1:
 207                     line = (yield p.ss.readline())
 208                     if self.debuglevel > 0:
 209                         print "reading from proxy", repr(line)
 210                     res += line
 211                     if line in ['\n', '\r\n']:
 212                         break
 213             except StopIteration:
 214                 if self.debuglevel > 0:
 215                     print "proxy read stopiter"
 216                 # TODO: close connection
 217             except:
 218                 if self.debuglevel > 0:
 219                     print 'proxy read error', \
 220                           (traceback and traceback.print_exc() or None)
 221                 # TODO: close connection
 222
 223             f = StringIO(res)
 224
 225             # Examine the headers and look for a Connection directive
 226             respheaders = mimetools.Message(f, 0)
 227             if self.debuglevel > 0:
 228                 print "response headers", str(respheaders)
 229             remote = self.client.remote
 230             rcooks = httpd.process_cookies(respheaders, remote, "Set-Cookie", False)
 231             rcooks['session'] = self.hr.response_cookies['session'].value # nooo
 232             rcooks['session']['expires'] = \
 233                     self.hr.response_cookies['session']['expires']
 234             self.hr.response_cookies = rcooks
 235             if self.debuglevel > 0:
 236                 print "rcooks", str(rcooks)
 237
 238             # override connection: keep-alive hack
 239             #responseline = responseline.split(" ")
 240             #print "responseline:", responseline
 241             #if responseline[1] != "200":
 242             #    respheaders['Connection'] = 'close'
 243
 244             # send all but Set-Cookie headers
 245             del respheaders['Set-Cookie'] # being replaced
 246             yield self.client.writeMessage(str(respheaders))
 247
 248             # now replacement cookies
 249             for k, v in rcooks.items():
 250                 val = v.output()
 251                 yield self.client.writeMessage(val+"\r\n")
 252
 253             # check connection for "closed" header
 254             if keepalive:
 255                 conntype = respheaders.get('Connection', "")
 256                 if conntype.lower() == 'close':
 257                     self.hr.close_connection = 1
 258                 elif (conntype.lower() == 'keep-alive' and
 259                       self.hr.protocol_version >= "HTTP/1.1"):
 260                     self.hr.close_connection = 0
 261
 262             # write rest of data
 263             if self.debuglevel > 0:
 264                 print "writing to client body"
 265             yield self.client.writeMessage("\r\n")
 266
 267             if respheaders.has_key('content-length'):
 268                 max_chunk_size = 10*1024*1024
 269                 size_remaining = int(respheaders["content-length"])
 270                 while size_remaining:
 271                     chunk_size = min(size_remaining, max_chunk_size)
 272                     data = (yield p.ss.read(chunk_size))
 273                     if self.debuglevel > 0:
 274                         print "reading from proxy expecting", \
 275                                     size_remaining, repr(data)
 276                     yield self.client.writeMessage(data)
 277                     size_remaining -= len(data)
 278             else:
 279                 while True:
 280                     #data = p.read()
 281                     try:
 282                         data = (yield p.ss.read(1024))
 283                     except httpd.ConnectionClosed:
 284                         break
 285                     if self.debuglevel > 0:
 286                         print "reading from proxy", repr(data)
 287                     if data == '':
 288                         break
 289                     yield self.client.writeMessage(data)
 290
 291             if not keepalive: #self.hr.close_connection:
 292                 if self.debuglevel > 0:
 293                     print 'proxy wants client to close_connection'
 294                 try:
 295                     yield self.client.connectionClosed()
 296                     p.serving = False
 297                     raise httpd.ConnectionClosed
 298                 except httpd.ConnectionClosed:
 299                     if self.debuglevel > 0:
 300                         print 'close_connection done'
 301                     pass
 302
 303             p.serving = False
 304         except httpd.ConnectionClosed:
 305             # whoops, remote end has died: remove client and
 306             # remove proxy session, we cannot do anything else,
 307             # there's nothing there to talk to.
 308             self.client.removeConnection()
 309             self.proxies.pop(session)
 310
 311         except:
 312             if self.debuglevel > 0:
 313                 print traceback.print_exc()
 314
 315         p.serving = False
 316         raise StopIteration
 317