a simple CONNECT http proxy

http_connect_proxy.py

# SPDX-License-Identifier: MIT

import asyncio
import socket


async def pipe(reader, writer):
    try:
        while not reader.at_eof():
            writer.write(await reader.read(2048))
    except ConnectionResetError:
        pass
    finally:
        writer.close()


def write_response(writer, resp):
    writer.write(
        "\r\n".join(
            [
                resp,
                "",
                "",
            ]
        ).encode()
    )


async def handle_client(local_reader, local_writer):
    try:
        request = await local_reader.read(2048)
        lines = request.split(b"\r\n")
        if not lines[0].startswith(b"CONNECT "):
            print("bad request", request.decode())
            write_response(local_writer, "HTTP/1.1 407 Only CONNECT allowed")
            return
        host, port = lines[0].split(b" ")[1].split(b":")
        try:
            remote_reader, remote_writer = await asyncio.open_connection(
                host.decode(), int(port)
            )
        except socket.gaierror:
            print("failed to relay to", host, port)
            write_response(local_writer, "HTTP/1.1 404 Not Found")
            return

        print("relaying to", host, port)
        write_response(local_writer, "HTTP/1.1 200 Connection established")
        pipe1 = pipe(local_reader, remote_writer)
        pipe2 = pipe(remote_reader, local_writer)
        await asyncio.gather(pipe1, pipe2)

    finally:
        local_writer.close()


loop = asyncio.get_event_loop()
coro = asyncio.start_server(handle_client, "127.0.0.1", 1080)
server = loop.run_until_complete(coro)

print("Serving on {}".format(server.sockets[0].getsockname()))
try:
    loop.run_forever()
except KeyboardInterrupt:
    pass

server.close()
loop.run_until_complete(server.wait_closed())
loop.close()