cheatproxy.py

   1 #!/usr/bin/env python
   2
   3 """A basic HTTP proxy server that intercepts communication with
   4 BitTorrent trackers and optionally spoofs the amount of data
   5 uploaded. Free from artificial colours and preservatives. Web 2.0
   6 compatible."""
   7
   8 import getopt
   9 import logging
  10 import select
  11 import socket
  12 import sys
  13 import urlparse
  14
  15 import BaseHTTPServer
  16 import SocketServer
  17
  18 import cheatbt
  19
  20 class CheatHandler(BaseHTTPServer.BaseHTTPRequestHandler):
  21     """Used by HTTPServer to handle HTTP requests"""
  22
  23     mappings = {}
  24
  25     def do_GET(self):
  26         """Called by BaseHTTPRequestHandler when a GET request is
  27         received from a client."""
  28
  29         cheatpath = cheatbt.cheat(self.path, CheatHandler.mappings)
  30
  31         logger = logging.getLogger("cheatproxy")
  32         logger.info(cheatpath)
  33
  34         (scheme, netloc, path, params, query, fragment) = \
  35             urlparse.urlparse(cheatpath, 'http')
  36
  37         # TODO: https support.
  38         if scheme != 'http' or fragment or not netloc:
  39             self.send_error(501)
  40             return
  41
  42         soc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  43
  44         try:
  45             if self._connect_to(netloc, soc):
  46                 request = urlparse.urlunparse(('', '', path, params, query, ''))
  47                 soc.send("%s %s %s\r\n" % (self.command, request,
  48                                            self.request_version))
  49                 self.headers['Connection'] = 'close'
  50                 del self.headers['Proxy-Connection']
  51                 # This is naughty. But rfc822.Message, which self.headers is a
  52                 # subclass of, insists on converting headers to lowercase when
  53                 # accessed conventionally (i.e. as a dict).
  54                 for header in self.headers.headers:
  55                     header = header.strip() + '\r\n'
  56                     soc.send(header)
  57                     logger.debug(repr(header))
  58                 soc.send("\r\n")
  59                 self._read_write(soc)
  60         finally:
  61             soc.close()
  62             self.connection.close()
  63
  64     def _connect_to(self, netloc, soc):
  65         """Attempt to establish a connection to the tracker."""
  66
  67         i = netloc.find(':')
  68         if i >= 0:
  69             host_port = (netloc[:i], int(netloc[i+1:]))
  70         else:
  71             host_port = (netloc, 80)
  72
  73         try:
  74             soc.connect(host_port)
  75         except socket.error:
  76             return False
  77
  78         return True
  79
  80     def _read_write(self, soc, max_idling=20):
  81         """Pass data between the remote server and the client. I
  82         think."""
  83
  84         logger = logging.getLogger("cheatproxy")
  85
  86         rlist = [self.connection, soc]
  87         wlist = []
  88         count = 0
  89         while True:
  90             count += 1
  91             (ins, _, exs) = select.select(rlist, wlist, rlist, 3)
  92             if exs:
  93                 break
  94             if ins:
  95                 for i in ins:
  96                     if i is soc:
  97                         out = self.connection
  98                     else:
  99                         out = soc
 100                     try:
 101                         data = i.recv(8192)
 102                         if data:
 103                             out.send(data)
 104                             count = 0
 105                     except socket.error, msg:
 106                         logger.info(msg)
 107             else:
 108                 pass
 109             if count == max_idling:
 110                 break
 111
 112 class CheatProxy(SocketServer.ThreadingMixIn, BaseHTTPServer.HTTPServer):
 113     def __init__(self, host, port):
 114         BaseHTTPServer.HTTPServer.__init__(self, (host, port),
 115                 CheatHandler)
 116
 117 def usage():
 118     """Prints usage information and exits."""
 119
 120     print """
 121 usage: %s [-b host] [-p port] [-f file] [-v] [-d] [-h]
 122
 123   -b host  IP or hostname to bind to. Default is localhost.
 124   -p port  Port to listen on. Default is 8000.
 125   -f file  Mappings file.
 126   -v       Verbose output.
 127   -d       Debug output.
 128   -h       What you're reading.
 129 """ % sys.argv[0]
 130     sys.exit(1)
 131
 132 def main():
 133     host = "localhost"
 134     port = 8000
 135
 136     rootlogger = logging.getLogger("")
 137     ch = logging.StreamHandler()
 138     ch.setFormatter(
 139         logging.Formatter("%(asctime)s:%(name)s:%(levelname)s %(message)s"))
 140     rootlogger.addHandler(ch)
 141
 142     try:
 143         opts, _ = getopt.getopt(sys.argv[1:], "b:f:p:hvd")
 144     except getopt.GetoptError:
 145         usage()
 146
 147     for opt, val in opts:
 148         if opt == "-b":
 149             host = val
 150         if opt == "-p":
 151             port = int(val)
 152         if opt == "-f":
 153             CheatHandler.mappings = cheatbt.load_mappings(val)
 154         if opt == "-v":
 155             rootlogger.setLevel(logging.INFO)
 156         if opt == "-d":
 157             rootlogger.setLevel(logging.DEBUG)
 158         if opt == "-h":
 159             usage()
 160
 161     httpd = CheatProxy(host, port)
 162
 163     logger = logging.getLogger("cheatproxy")
 164     logger.info("listening on %s:%d" % (host, port))
 165
 166     httpd.serve_forever()
 167
 168 if __name__ == "__main__":
 169     try:
 170         main()
 171     except KeyboardInterrupt:
 172         logging.shutdown()
 173         sys.exit()