import socket
import concurrent.futures
import select
import logging

SOCKS_VERSION = 5
HOST = "0.0.0.0"
PORT = 5678
WORKER = 64


class Proxy:

    def __init__(self):
        self.username = "tele"
        self.password = "proxy"
        self.BUFF_SIZE = 8192  # Default 4096
        self.active_connections = 0

    def handle_client(self, connection):
        """
        Handle incoming client connections.
        """
        remote = None
        address = None
        try:
            version, nmethods = connection.recv(2)
            methods = self.get_available_methods(nmethods, connection)

            if 2 not in set(methods):
                connection.close()
                return

            connection.sendall(bytes([SOCKS_VERSION, 2]))

            if not self.verify_credentials(connection):
                return

            data = connection.recv(4)
            if len(data) != 4:
                connection.close()
                return

            version, cmd, _, address_type = data

            if address_type == 1:
                address = socket.inet_ntoa(connection.recv(4))
            elif address_type == 3:
                domain_length = connection.recv(1)[0]
                address = connection.recv(domain_length)
                address = socket.gethostbyname(address)

            port = int.from_bytes(connection.recv(2), 'big', signed=False)

            if address is None:
                connection.close()
                return

            remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            remote.connect((address, port))
            bind_address = remote.getsockname()

            reply = self.generate_reply(bind_address)
            connection.sendall(reply)

            if reply[1] == 0 and cmd == 1:
                self.exchange_loop(connection, remote)

        except KeyboardInterrupt:
            pass
        except ConnectionResetError:
            pass
        except Exception as e:
            logging.error("{}".format(e))
        finally:
            try:
                connection.close()
                if remote:
                    remote.close()
                self.active_connections -= 1
                logging.info("Closed current connection. Current active connections: {}".format(
                    self.active_connections))
            except Exception as e:
                logging.error("Error closing connection: {}".format(e))

    def exchange_loop(self, client, remote):
        """
        Exchange data between client and remote server.
        """
        while True:
            r, _, _ = select.select([client, remote], [], [])
            for sock in r:
                data = sock.recv(self.BUFF_SIZE)
                if not data:
                    return
                if sock is client:
                    remote.sendall(data)
                else:
                    client.sendall(data)

    def generate_reply(self, bind_address):
        """
        Generate SOCKS5 reply message.
        """
        if bind_address[0] is None:
            return b''

        addr_bytes = socket.inet_aton(bind_address[0])
        port_bytes = bind_address[1].to_bytes(2, 'big')
        return b''.join([
            SOCKS_VERSION.to_bytes(1, 'big'),
            int(0).to_bytes(1, 'big'),
            int(0).to_bytes(1, 'big'),
            int(1).to_bytes(1, 'big'), addr_bytes, port_bytes
        ])

    def verify_credentials(self, connection):
        """
        Verify client credentials.
        """
        version = ord(connection.recv(1))
        username_len = ord(connection.recv(1))
        username = connection.recv(username_len).decode('utf-8')
        password_len = ord(connection.recv(1))
        password = connection.recv(password_len).decode('utf-8')

        if username == self.username and password == self.password:
            response = bytes([version, 0])
            connection.sendall(response)
            return True
        else:
            response = bytes([version, 0xFF])
            connection.sendall(response)
            return False

    def get_available_methods(self, nmethods, connection):
        """
        Get available SOCKS5 authentication methods.
        """
        methods = []
        for i in range(nmethods):
            methods.append(ord(connection.recv(1)))
        return methods

    def run(self):
        """
        Run the SOCKS5 proxy server.
        """
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            s.bind((HOST, PORT))
            s.listen()

            logging.info("Socks5 proxy server is running on {}:{}".format(
                HOST, PORT))

            with concurrent.futures.ThreadPoolExecutor(
                    max_workers=WORKER) as executor:
                try:
                    while True:
                        conn, addr = s.accept()
                        self.active_connections += 1
                        logging.info(
                            "New connection from {}. Active connections: {}".
                            format(addr, self.active_connections))
                        executor.submit(self.handle_client, conn)
                except KeyboardInterrupt:
                    logging.info("Closing server socket.")
                    conn.send(b'')
                    s.shutdown(socket.SHUT_RDWR)
                    s.close()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s - %(levelname)s - %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S')
    proxy = Proxy()
    proxy.run()
