from http_parser.http import HttpStream
from http_parser.reader import SocketReader
import requests
import socketserver
import re
import argparse

def response2http(resp):
    body = resp.raw.read()
    resp.close()

    s = f'HTTP/1.1 {resp.status_code} {resp.reason}\r\n'
    for key in resp.headers:
        s = s + f'{key}: {resp.headers[key]}\r\n'
    s = s + f'Content-Length: {len(body)}\r\n'
    s = s + '\r\n'

    return bytes(s, encoding='utf-8') + body

class ProxyHandler(socketserver.BaseRequestHandler):
    def handle(self):
        reader = SocketReader(self.request)
        stream = HttpStream(reader)

        url = stream.url()
        if not url.startswith("http"):
            url = "https://" + url
        else:
            url = re.sub(r"^http", "https", stream.url(), count=1)

        print(url)
        # Disable compression by default, because some HTTP clients do
        # not support compression.
        headers = stream.headers()
        if 'Accept-Encoding' not in headers:
            headers['Accept-Encoding'] = 'identity'

        r = requests.get(url, headers=headers, stream=True)

        respBytes = response2http(r)

        sent = 0
        while sent < len(respBytes):
            sent += self.request.send(respBytes[sent:])
            if sent == 0:
                break

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-l", "--listen", help="Listen IP", action="store", default="0.0.0.0")
    parser.add_argument("-p", "--port", help="Proxy Port", action="store", type=int, default=8080)
    args = parser.parse_args()

    with socketserver.ForkingTCPServer((args.listen, args.port), ProxyHandler) as server:
        server.serve_forever()