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