#!/usr/bin/env python3
"""Logging SOCKS5 server for the loopback lab: no authentication, CONNECT only.

It prints one line per request (address type and destination), so a run shows whether a client
spoke SOCKS at all and whether it sent a hostname or an IP address. A client that does not start
with the SOCKS version byte 0x05 (for example one that sends an HTTP request to this port) is
logged and disconnected. Bind it to 127.0.0.1 only; it is a test fixture, not a production proxy.

    python3 socks5_server.py 1080
"""
import select
import socket
import socketserver
import struct
import sys

ATYP = {1: "IPv4", 3: "DOMAIN", 4: "IPv6"}


class Handler(socketserver.BaseRequestHandler):
    def read(self, size):
        data = b""
        while len(data) < size:
            chunk = self.request.recv(size - len(data))
            if not chunk:
                raise ConnectionError("client closed the connection")
            data += chunk
        return data

    def handle(self):
        try:
            greeting = self.read(2)
            if greeting[0] != 5:
                more = self.request.recv(64)
                print(f"NOT_SOCKS5 first_bytes={(greeting + more)[:24]!r}", flush=True)
                return
            self.read(greeting[1])
            self.request.sendall(b"\x05\x00")  # no authentication
            _version, command, _reserved, atyp = self.read(4)
            if atyp == 1:
                host = socket.inet_ntoa(self.read(4))
            elif atyp == 3:
                host = self.read(self.read(1)[0]).decode("ascii", "replace")
            else:
                host = socket.inet_ntop(socket.AF_INET6, self.read(16))
            port = struct.unpack("!H", self.read(2))[0]
            print(f"CONNECT atyp={ATYP.get(atyp, atyp)} dest={host}:{port}", flush=True)
            if command != 1:
                self.request.sendall(b"\x05\x07\x00\x01" + bytes(6))  # command not supported
                return
            try:
                upstream = socket.create_connection((host, port), timeout=10)
            except OSError:
                self.request.sendall(b"\x05\x05\x00\x01" + bytes(6))  # connection refused
                return
            self.request.sendall(b"\x05\x00\x00\x01" + bytes(6))
            with upstream:
                pair = {self.request: upstream, upstream: self.request}
                while True:
                    readable, _, _ = select.select(list(pair), [], [], 60)
                    if not readable:
                        return
                    for source in readable:
                        data = source.recv(65536)
                        if not data:
                            return
                        pair[source].sendall(data)
        except (ConnectionError, OSError):
            return


class Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 1080
    with Server(("127.0.0.1", port), Handler) as server:
        print(f"SOCKS5 server listening on 127.0.0.1:{port}", flush=True)
        server.serve_forever()