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     def close(self):
  77         """Close the connection to the HTTP server."""
  78         if self.sock:
  79             self.sock.close()   # close it manually... there may be other refs
  80             self.sock = None
  81
  82     def send(self, str):
  83         """Send `str' to the server."""
  84         if self.sock is None:
  85             if self.auto_open:
  86                 self.connect()
  87             else:
  88                 raise NotConnected()
  89
  90         # send the data to the server. if we get a broken pipe, then close
  91         # the socket. we want to reconnect when somebody tries to send again.
  92         #
  93         # NOTE: we DO propagate the error, though, because we cannot simply
  94         #       ignore the error... the caller will know if they can retry.
  95         if self.debuglevel > 0:
  96             print "send:", repr(str)
  97         try:
  98             yield multitask.send(self.sock, str)
  99         except socket.error, v:
 100             if v[0] == 32:      # Broken pipe
 101                 self.close()
 102             raise
 103
 104     def read(self, num_bytes=1024):
 105         data = (yield multitask.recv(self.sock, num_bytes))
 106
 107
 108 class ProxyServerRequestHandler(object):
 109
 110     """Simple HTTP request handler with GET and HEAD commands.
 111
 112     This serves files from the current directory and any of its
 113     subdirectories.  The MIME type for files is determined by
 114     calling the .guess_type() method.
 115
 116     The GET and HEAD requests are identical except that the HEAD
 117     request omits the actual contents of the file.
 118
 119     """
 120
 121     server_version = "SimpleHTTP/" + __version__
 122
 123     def on_query(self, client, reqtype, *args):
 124         """Serve a request."""
 125         self.client = client
 126         self.hr = args[0]
 127         if not hasattr(self.client, "proxy"):
 128             self.client.proxy = ProxyConnection()
 129             self.client.proxy.connect()
 130
 131         multitask.add(self.proxy_relay(reqtype))
 132
 133         return True
 134
 135     def onPOST(self, client, *args):
 136         """Serve a POST request."""
 137         return self.on_query(client, "POST", *args)
 138
 139     def onGET(self, client, *args):
 140         """Serve a GET request."""
 141         return self.on_query(client, "GET", *args)
 142
 143     def proxy_relay(self, reqtype):
 144
 145         p = self.client.proxy
 146
 147         # send command
 148         req = "%s %s %s\n" % (reqtype, self.hr.path, self.hr.request_version)
 149         print "req", req
 150         yield multitask.send(p.sock, req)
 151
 152         # send headers
 153         hdrs = str(self.hr.headers)
 154         print "hdrs", hdrs
 155         yield multitask.send(p.sock, hdrs)
 156         yield multitask.send(p.sock, "\n")
 157
 158         # now content
 159         if self.hr.headers.has_key('content-length'):
 160             max_chunk_size = 10*1024*1024
 161             size_remaining = int(self.hr.headers["content-length"])
 162             L = []
 163             print "size_remaining", size_remaining
 164             while size_remaining:
 165                 chunk_size = min(size_remaining, max_chunk_size)
 166                 data = self.hr.rfile.read(chunk_size)
 167                 yield multitask.send(p.sock, data)
 168                 size_remaining -= len(L[-1])
 169
 170         # now read response and write back
 171         res = ''
 172         while True:
 173             #data = p.read()
 174             data = (yield multitask.recv(p.sock, 1024))
 175             print "reading from proxy", repr(data)
 176             if data == '':
 177                 break
 178             res += data
 179
 180         f = StringIO(res)
 181         requestline = f.readline()
 182         yield self.client.writeMessage(requestline)
 183
 184         # Examine the headers and look for a Connection directive
 185         respheaders = mimetools.Message(f, 0)
 186         print "response headers", str(respheaders)
 187         remote = self.client.remote
 188         rcooks = httpd.process_cookies(respheaders, remote, "Set-Cookie", False)
 189         rcooks['session'] = self.hr.response_cookies['session'].value # nooo
 190         rcooks['session']['expires'] = \
 191                 self.hr.response_cookies['session']['expires']
 192         self.hr.response_cookies = rcooks
 193         print "rcooks", str(rcooks)
 194
 195         # send all but Set-Cookie headers
 196         del respheaders['Set-Cookie'] # being replaced
 197         yield self.client.writeMessage(str(respheaders))
 198
 199         # now replacement cookies
 200         for k, v in rcooks.items():
 201             val = v.output()
 202             yield self.client.writeMessage(val+"\r\n")
 203
 204         # check connection for "closed" header
 205         conntype = respheaders.get('Connection', "")
 206         if conntype.lower() == 'close':
 207             self.hr.close_connection = 1
 208         elif (conntype.lower() == 'keep-alive' and
 209               self.hr.protocol_version >= "HTTP/1.1"):
 210             self.hr.close_connection = 0
 211
 212         # write rest of data
 213         print "writing to client body"
 214         yield self.client.writeMessage("\r\n")
 215         yield self.client.writeMessage(f.read())
 216         if self.hr.close_connection:
 217             try:
 218                 yield self.client.connectionClosed()
 219             except httpd.ConnectionClosed:
 220                 print 'close_connection done'
 221                 pass
 222
 223         raise StopIteration
 224