#!/usr/bin/env python3
"""Loopback-only teaching CONNECT and SOCKS5 CONNECT subset, no encryption."""
import argparse
import http.client
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
import json
import select
import socket
import socketserver
import struct
import threading


def exact(sock,n):
    result=bytearray()
    while len(result)<n:
        part=sock.recv(n-len(result))
        if not part:raise EOFError('short proxy request')
        result.extend(part)
    return bytes(result)


def relay(left,right):
    readers=[left,right]
    while readers:
        ready,_,_=select.select(readers,[],[],5)
        if not ready:raise TimeoutError('lab relay idle timeout')
        for source in ready:
            target=right if source is left else left
            data=source.recv(4096)
            if data:target.sendall(data)
            else:
                readers.remove(source)
                try:target.shutdown(socket.SHUT_WR)
                except OSError:pass  # The peer may already have closed/reset its side.


class Origin(BaseHTTPRequestHandler):
    protocol_version='HTTP/1.1'
    def do_GET(self):
        body=json.dumps({'origin_peer':list(self.client_address),'path':self.path}).encode()
        self.send_response(200);self.send_header('Content-Length',str(len(body)))
        self.end_headers();self.wfile.write(body)
    def log_message(self,*args):pass


class Proxy(socketserver.BaseRequestHandler):
    def handle(self):
        client=self.request;client.settimeout(5)
        host=None;port=None;mode=None
        if self.server.mode=='connect':
            try:
                header=bytearray()
                while not header.endswith(b'\r\n\r\n'):
                    header.extend(exact(client,1))
                    if len(header)>8192:raise ValueError('header exceeds limit')
                method,target,version=header.split(b'\r\n',1)[0].decode('ascii').split()
                if method!='CONNECT' or version!='HTTP/1.1':raise ValueError('CONNECT subset only')
                host,port_text=target.rsplit(':',1);port=int(port_text);mode='HTTP CONNECT'
                if not 1<=port<=65535:raise ValueError('bad port')
            except (ValueError,EOFError,OSError):
                try:client.sendall(b'HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n')
                except OSError:pass
                return
        else:
            version,count=exact(client,2)
            if version!=5 or count==0:raise ValueError('bad SOCKS greeting')
            methods=exact(client,count)
            if 0 not in methods:
                client.sendall(b'\x05\xff');return
            client.sendall(b'\x05\x00')
            version,command,reserved,atyp=exact(client,4)
            if (version,reserved)!=(5,0):return
            if atyp==1:host=socket.inet_ntop(socket.AF_INET,exact(client,4));mode='SOCKS IP'
            elif atyp==3:host=exact(client,exact(client,1)[0]).decode('ascii');mode='SOCKS domain'
            elif atyp==4:host=socket.inet_ntop(socket.AF_INET6,exact(client,16))
            else:
                client.sendall(b'\x05\x08\x00\x01'+b'\x00'*6);return
            port=struct.unpack('!H',exact(client,2))[0]
            # Consume known address formats before replying and closing, so the
            # queued request bytes do not cause a reset that obscures the reply.
            if command!=1:
                client.sendall(b'\x05\x07\x00\x01'+b'\x00'*6);return
            if atyp==4:
                client.sendall(b'\x05\x08\x00\x01'+b'\x00'*6);return
        # Static lesson.test alias at the proxy. No external DNS is queried.
        if host not in ('lesson.test','127.0.0.1') or port!=self.server.origin_port:
            if self.server.mode=='connect':client.sendall(b'HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\n\r\n')
            else:client.sendall(b'\x05\x02\x00\x01'+b'\x00'*6)
            return
        with socket.create_connection(('127.0.0.1',port),timeout=5) as upstream:
            event={'mode':mode,'requested_host':host,'requested_port':port,
                'resolution':'static alias at proxy' if host=='lesson.test' else 'IP already provided',
                'proxy_client_peer':list(client.getpeername()),'proxy_origin_source':list(upstream.getsockname())}
            with self.server.events_lock:self.server.events.append(event)
            if self.server.mode=='connect':client.sendall(b'HTTP/1.1 200 Connection Established\r\n\r\n')
            else:
                address,source_port=upstream.getsockname()
                client.sendall(b'\x05\x00\x00\x01'+socket.inet_aton(address)+struct.pack('!H',source_port))
            relay(client,upstream)


class Server(socketserver.ThreadingTCPServer):
    daemon_threads=True


def http_request(sock):
    sock.sendall(b'GET /lesson HTTP/1.1\r\nHost: lesson.test\r\nConnection: close\r\n\r\n')
    response=http.client.HTTPResponse(sock);response.begin();body=response.read()
    return {'status':response.status,'response':json.loads(body)}


def run_case(server,use_domain):
    with socket.create_connection(server.server_address,timeout=5) as client:
        endpoint=list(client.getsockname())
        host='lesson.test' if use_domain else '127.0.0.1'
        if server.mode=='connect':
            client.sendall(f'CONNECT {host}:{server.origin_port} HTTP/1.1\r\nHost: {host}\r\n\r\n'.encode())
            header=bytearray()
            while not header.endswith(b'\r\n\r\n'):header.extend(exact(client,1))
            if not header.startswith(b'HTTP/1.1 200 '):raise RuntimeError('CONNECT denied')
        else:
            client.sendall(b'\x05\x01\x00')
            if exact(client,2)!=b'\x05\x00':raise RuntimeError('auth negotiation failed')
            address=(b'\x03'+bytes([len(host)])+host.encode()) if use_domain else b'\x01'+socket.inet_aton(host)
            client.sendall(b'\x05\x01\x00'+address+struct.pack('!H',server.origin_port))
            reply=exact(client,10)
            if reply[:4]!=b'\x05\x00\x00\x01':raise RuntimeError('SOCKS request denied')
        result=http_request(client);result['client_source']=endpoint
        return result


def demo():
    with ThreadingHTTPServer(('127.0.0.1',0),Origin) as origin,Server(('127.0.0.1',0),Proxy) as proxy:
        proxy.origin_port=origin.server_port;proxy.events=[];proxy.events_lock=threading.Lock()
        threads=[threading.Thread(target=x.serve_forever) for x in (origin,proxy)]
        for t in threads:t.start()
        try:
            results=[]
            proxy.mode='connect';results.append(run_case(proxy,True))
            proxy.mode='socks';results.append(run_case(proxy,True));results.append(run_case(proxy,False))
            # Cross-check the distinct TCP endpoints reported by proxy and origin.
            for result,event in zip(results,proxy.events):
                assert result['response']['origin_peer']==event['proxy_origin_source']
                assert result['client_source']==event['proxy_client_peer']
            return {'kind':'real loopback forwarding; static name fixture, not Internet DNS',
                    'results':results,'proxy_events':proxy.events,
                    'subset':'CONNECT and SOCKS5 CONNECT only; no UDP ASSOCIATE, TLS, IPv6 or general destination access.'}
        finally:
            proxy.shutdown();origin.shutdown()
            for t in threads:t.join(timeout=3)


if __name__=='__main__':print(json.dumps(demo(),ensure_ascii=False,indent=2))