http/tiny-http-proxy

   1 #!/usr/bin/python
   2 #
   3 # A Tiny HTTP Proxy.
   4 #
   5 # TODO:
   6 #
   7 #   - Doesn't understand chunked data
   8 #   - Try: www.wikipedia.org
   9 #
  10 # Luiz Capitulino <lcapitulino@gmail.com>
  11
  12 import socket
  13 from sys import argv, exit, stderr
  14
  15 class HttpMessage:
  16     def get_server(self):
  17         uri = self.header.split()[1]
  18         return uri[uri.find('//') + 2:uri.rfind('/')]
  19
  20     def conn_close(self):
  21         for line in self.header.split('\n'):
  22             if line.startswith('Connection') and line.find('close') != -1:
  23                 return True
  24         return False
  25
  26     def __repr__(self):
  27         return self.message
  28
  29     def __init__(self, header, body):
  30         self.header = header
  31         self.body = body
  32         self.message = self.header
  33         if self.body:
  34             self.message += self.body
  35         print self.header
  36
  37 class TinyHttpProxy:
  38     def __read_header(self, sock):
  39         header = ''
  40         in_end = False
  41         while True:
  42             b = sock.recv(1)
  43             header += b
  44             if not b:
  45                 return
  46             if b == '\r':
  47                 continue
  48             elif b != '\n':
  49                 in_end = False
  50                 continue
  51             elif in_end:
  52                 return header
  53             else:
  54                 in_end = True
  55
  56     def __get_content_length(self, header):
  57         for line in header.split('\n'):
  58             if line.startswith('Content-Length'):
  59                 return int(line[line.find(' '):])
  60         return 0
  61
  62     def __read_body(self, header, sock):
  63         buf = ''
  64         len = self.__get_content_length(header)
  65         for i in range(len):
  66             buf += sock.recv(1)
  67         return buf
  68
  69     def __read_message(self, sock):
  70         header = self.__read_header(sock)
  71         if not header:
  72             return
  73         body = self.__read_body(header, sock)
  74         return HttpMessage(header, body)
  75
  76     def __create_proxy_sock(self, domain, port):
  77         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  78         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  79         s.bind((domain, port))
  80         s.listen(1)
  81         return s
  82
  83     def __server_connect(self, server, port):
  84         ssock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  85         ssock.connect((server, port))
  86         return ssock
  87
  88     def main_loop(self):
  89         while True:
  90             ssock = None
  91             (csock, addr) = self.sock.accept()
  92             req = self.__read_message(csock)
  93             while True:
  94                 if not ssock:
  95                     ssock = self.__server_connect(req.get_server(), 80)
  96                 ssock.sendall(req.message)
  97                 resp = self.__read_message(ssock)
  98                 if not resp:
  99                     ssock.close()
 100                     ssock = None
 101                     continue
 102                 csock.sendall(resp.message)
 103                 req = self.__read_message(csock)
 104                 if resp.conn_close():
 105                     ssock.close()
 106                     csock.close()
 107                     break
 108
 109     def __init__(self, domain, port):
 110         self.sock = self.__create_proxy_sock(domain, port)
 111
 112 def main():
 113     if len(argv) != 3:
 114         print 'tiny-http-proxy <domain> <port>'
 115         exit(1)
 116
 117     server = TinyHttpProxy(argv[1], int(argv[2]))
 118     server.main_loop()
 119
 120 if __name__ == '__main__':
 121     main()