# -*- coding: utf-8 -*-
import socket
import struct
import select
import sys
import logging
import threading
from socketserver import StreamRequestHandler
from socketserver import ThreadingTCPServer as RawThreadingTCPServer
import http.client

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
AUTH_USERNAME = "YOUR_PROXY_LOGIN"
AUTH_PASSWORD = "YOUR_PROXY_PASSWORD"


#
# Constants
#
'''Version of the protocol'''
# PROTOCOL VERSION 5
VER = 5
'''Method constants'''
# '00' NO AUTHENTICATION REQUIRED
M_NOAUTH = b'\x00'
# '02' USER/PASSWORD
M_USER_PASS = b'\x02'
# 'FF' NO ACCEPTABLE METHODS
M_NOTAVAILABLE = b'\xff'
'''Command constants'''
# CONNECT '01'
CMD_CONNECT = 1
'''Address type constants'''
# IP V4 address '01'
CMD_UDP_ASSOCIATE = 3
ATYP_IPV4 = 1
# DOMAINNAME '03'
ATYP_DOMAINNAME = 3
ATYP_IPV6 = 4

REMOTE_TIMEOUT = 10

def debugger_is_active() -> bool:
    """Return if the debugger is currently active"""
    return hasattr(sys, 'gettrace') and sys.gettrace() is not None

class ThreadingTCPServer(RawThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True
    # if debugger_is_active():
    #     request_queue_size = 1 # For debugging


class SocksException(Exception):
    pass


class Socks5Handler(StreamRequestHandler):

    def handle(self):
        try:
            self.socks_auth()  # Perform SOCKS authentication

            ver, cmd, _, addr_type = self.rfile.read(4)
            if ver != VER:
                raise SocksException("invalid version " + str(ver))

            remote_addr, remote_port = self.remote_info(addr_type)
            logging.info("%s:%s", remote_addr, remote_port)

            if cmd == CMD_UDP_ASSOCIATE:  # Handle UDP ASSOCIATE
                # Create a UDP socket and send its details back to the client
                udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                udp_sock.bind(('0.0.0.0', 0))
                udp_addr, udp_port = udp_sock.getsockname()

                reply = b'\x05\x00\x00\x01'
                reply += socket.inet_aton(udp_addr) + struct.pack('>H', udp_port)
                self.wfile.write(reply)

                self.handle_udp_associate(udp_sock)  # Handle UDP traffic

            elif cmd == CMD_CONNECT:
                reply = b'\x05\x00\x00\x01'
                reply += socket.inet_aton('0.0.0.0') + struct.pack('>H', 2222)
                self.wfile.write(reply)
                self.duplex_transport(remote_addr, remote_port)  # Perform duplex transport

        except SocksException as e:
            logger.error("SOCKS error: %r", e)
        except socket.error as e:
            logger.error("Socket error: %r", e)


    def duplex_transport(self, remote_addr, remote_port):
        """
        Perform duplex transport between the client and the target.
        """
        remote = socket.create_connection((remote_addr, remote_port), timeout=REMOTE_TIMEOUT)
        self.handle_forward_in_select(self.connection, remote)


    def handle_forward_in_select(self, client, remote):
        # type: (socket.socket, socket.socket) -> None
        try:
            inputs = [client, remote]
            while inputs:
                readable, _, _ = select.select(inputs, [], [])
                if client in readable:
                    data = client.recv(4096)
                    if not data:
                        break
                    remote.sendall(data)

                if remote in readable:
                    data = remote.recv(4096)
                    if not data:
                        break
                    client.sendall(data)
        except Exception as e:
            logger.error("select error:%r", e)

    def socks_user_auth(self):
        buf = self.rfile.read(2)
        if buf[0] != 0x01:
            raise SocksException("invalid version")

        username_length = buf[1]
        username = self.rfile.read(username_length).decode("utf-8")

        buf = self.rfile.read(1)
        password_length = buf[0]
        password = self.rfile.read(password_length).decode("utf-8")
        logger.debug("user auth: %s-%s", username, password)

        if username == AUTH_USERNAME and password == AUTH_PASSWORD:
            self.wfile.write(b'\x01\x00')
        else:
            self.wfile.write(b'\x01\x01')
            raise SocksException("socks user auth fail")

    def socks_auth(self):
        buf = self.rfile.read(2)
        if not len(buf) == 2 or buf[0] != VER:
            raise SocksException("invalid version")

        methods_length = buf[1]
        methods = self.rfile.read(methods_length)
        logger.debug("receive methods:%r", methods)
        if M_USER_PASS in methods:
            self.wfile.write(b'\x05\x02')
            self.socks_user_auth()
        else:
            self.wfile.write(b'\x05\x00')

    def remote_info(self, addr_type):
        if addr_type == ATYP_IPV4:
            logger.debug("client addr type:IPv4")
            packed_ip = self.rfile.read(4)
            remote_addr = socket.inet_ntoa(packed_ip)
            logger.debug(f"{remote_addr}:{packed_ip}")
        elif addr_type == ATYP_DOMAINNAME:
            logger.debug("client addr type:URL")
            addr_length = self.rfile.read(1)
            remote_addr = self.rfile.read(ord(addr_length))
            logger.debug(f"{remote_addr}")
        elif addr_type == ATYP_IPV6:
            logger.debug("client addr type:IPv6")
            packed_ip = self.rfile.read(16)
            remote_addr = socket.inet_ntop(socket.AF_INET6, packed_ip)
            logger.debug(f"{remote_addr}:{packed_ip}")
        else:
            raise SocksException("addr_type not supported.")

        if not isinstance(remote_addr, str):
            remote_addr = str(remote_addr, "utf-8")

        addr_port = self.rfile.read(2)
        remote_port = struct.unpack('>H', addr_port)[0]

        return remote_addr, remote_port


    def handle_udp_associate(self, udp_sock):
        """
        Handle UDP traffic between the client and the target.
        """
        client_sockets = {}  # Dictionary to store client addresses and their corresponding sockets

        try:
            while True:
                ready_sockets, _, _ = select.select([udp_sock] + list(client_sockets.values()), [], [])

                for sock in ready_sockets:
                    if sock == udp_sock:
                        # New UDP client connection
                        data, client_addr = udp_sock.recvfrom(4096)
                        if not data:
                            break

                        # Create a new socket for the client if it doesn't exist
                        if client_addr not in client_sockets:
                            client_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                            client_sockets[client_addr] = client_sock
                    else:
                        # Existing UDP client sending data
                        data, _ = sock.recvfrom(4096)
                        if not data:
                            # Client disconnected, remove the socket from the dictionary
                            client_sockets = {k: v for k, v in client_sockets.items() if v != sock}
                            sock.close()
                            continue

                        client_addr = next(k for k, v in client_sockets.items() if v == sock)

                    # Process the UDP data
                    rsv, frag, addr_type = struct.unpack('!HBB', data[:4])

                    dst_addr, dst_port, payload = self.parse_udp_header(addr_type, data[4:])

                    reply = self.relay_udp_payload(dst_addr, dst_port, payload)

                    # Construct the UDP reply header
                    reply_header = struct.pack("!HBB", 0, 0, addr_type)
                    if addr_type == ATYP_IPV4:
                        reply_header += socket.inet_aton(dst_addr)
                    elif addr_type == ATYP_DOMAINNAME:
                        encoded = dst_addr.encode() if isinstance(dst_addr, str) else dst_addr
                        reply_header += struct.pack("!B", len(encoded)) + encoded
                    elif addr_type == ATYP_IPV6:
                        reply_header += socket.inet_pton(socket.AF_INET6, dst_addr)
                    reply_header += struct.pack("!H", dst_port)

                    # Send the reply back to the client
                    reply = reply_header + reply
                    udp_sock.sendto(reply, client_addr)

        except Exception as e:
            logger.error("UDP associate error: %r", e)

        finally:
            # Close all client sockets
            for sock in client_sockets.values():
                sock.close()
            udp_sock.close()


    def relay_udp_payload(self, dst_addr, dst_port, payload):
        """Forward a UDP payload to dst_addr:dst_port and return the reply.

        Subclasses can override this to intercept specific traffic (e.g. DNS).
        """
        relay_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        relay_sock.settimeout(REMOTE_TIMEOUT)
        try:
            relay_sock.sendto(payload, (dst_addr, dst_port))
            reply, _ = relay_sock.recvfrom(4096)
        finally:
            relay_sock.close()
        return reply

    def parse_udp_header(self, addr_type, data):
        if addr_type == ATYP_IPV4:
            dst_addr = socket.inet_ntoa(data[:4])
            dst_port = struct.unpack("!H", data[4:6])[0]
            data = data[6:]
        elif addr_type == ATYP_DOMAINNAME:
            addr_len = data[0]
            dst_addr = data[1:1+addr_len].decode()
            dst_port = struct.unpack("!H", data[1+addr_len:1+addr_len+2])[0]
            data = data[1+addr_len+2:]
        elif addr_type == ATYP_IPV6:
            dst_addr = socket.inet_ntop(socket.AF_INET6, data[:16])
            dst_port = struct.unpack("!H", data[16:18])[0]
            data = data[18:]
        else:
            raise SocksException("Unsupported address type: " + str(addr_type))

        return dst_addr, dst_port, data




if __name__ == '__main__':
    server = ThreadingTCPServer(('127.0.0.1', 9050), Socks5Handler)
    server.serve_forever()

    # resolve_domain_doh("www.baidu.com")