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