#! /usr/bin/env python

from socket import create_connection
from select import select

try:  # Python < 3
    from SocketServer import TCPServer, BaseRequestHandler
    from urlparse import urlsplit
except ImportError:  # Python 3
    from socketserver import TCPServer, BaseRequestHandler
    from urllib.parse import urlsplit

def main(addr="localhost:0", limit=100):
    TCPServer.allow_reuse_address = True
    server = TCPServer(parse_addr(addr, defport=0), Handler)
    try:
        server.limit = limit
        print("Proxy server at http://" + format_addr(server.server_address))
        server.serve_forever()
    finally:
        server.server_close()

class Handler(BaseRequestHandler):
    def handle(self):
        reader = self.request.makefile("rb", 0)
        try:
            [method, target, _] = reader.readline().split(b" ", 2)
            if method != b"CONNECT":
                raise ValueError("Not a CONNECT request")
            target = target.decode("ascii")
            while reader.readline() not in {b"\r\n", b""}:
                pass
        finally:
            reader.close()
        
        conn = create_connection(parse_addr(target))
        try:
            print("Proxying connection to " + target)
            self.request.sendall(
                b"HTTP/1.1 200 Connection established\r\n"
                b"\r\n"
            )
            limit = self.server.limit
            total = 0
            while limit > 0:
                [ready, _, _] = select((self.request, conn), (), ())
                if self.request in ready:
                    data = self.request.recv(0x10000)
                    if not data:
                        break
                    conn.sendall(data)
                if conn in ready:
                    data = conn.recv(min(0x10000, limit))
                    if not data:
                        break
                    total += len(data)
                    limit -= len(data)
                    self.request.sendall(data)
            print("Forwarded {} B to client".format(total))
        finally:
            conn.close()

def format_addr(address):
    [address, port] = address
    address = "{}:{}".format(address, port)
    return address

def parse_addr(address, defport=None):
    address = urlsplit("//" + address)
    host = address.hostname
    if host is None:
        host = ""
    port = address.port
    if port is None:
        port = defport
    return (host, port)

if __name__ == "__main__":
    from argparse import ArgumentParser
    
    parser = ArgumentParser()
    [addr, limit] = main.__defaults__
    parser.add_argument("addr", nargs="?", default=addr)
    parser.add_argument("--limit", type=int, default=limit)
    args = parser.parse_args()
    main(args.addr, limit=args.limit)