#!/usr/bin/env python3
# Allow direct execution
import os
import sys
import threading
import unittest

import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import abc
import contextlib
import enum
import functools
import http.server
import json
import random
import socket
import struct
import time
from socketserver import (
    BaseRequestHandler,
    StreamRequestHandler,
    ThreadingTCPServer,
)

from test.helper import http_server_port
from yt_dlp.networking import Request
from yt_dlp.networking.exceptions import ProxyError, TransportError
from yt_dlp.socks import (
    SOCKS4_REPLY_VERSION,
    SOCKS4_VERSION,
    SOCKS5_USER_AUTH_SUCCESS,
    SOCKS5_USER_AUTH_VERSION,
    SOCKS5_VERSION,
    Socks5AddressType,
    Socks5Auth,
)

SOCKS5_USER_AUTH_FAILURE = 0x1


class Socks4CD(enum.IntEnum):
    REQUEST_GRANTED = 90
    REQUEST_REJECTED_OR_FAILED = 91
    REQUEST_REJECTED_CANNOT_CONNECT_TO_IDENTD = 92
    REQUEST_REJECTED_DIFFERENT_USERID = 93


class Socks5Reply(enum.IntEnum):
    SUCCEEDED = 0x0
    GENERAL_FAILURE = 0x1
    CONNECTION_NOT_ALLOWED = 0x2
    NETWORK_UNREACHABLE = 0x3
    HOST_UNREACHABLE = 0x4
    CONNECTION_REFUSED = 0x5
    TTL_EXPIRED = 0x6
    COMMAND_NOT_SUPPORTED = 0x7
    ADDRESS_TYPE_NOT_SUPPORTED = 0x8


class SocksTestRequestHandler(BaseRequestHandler):

    def __init__(self, *args, socks_info=None, **kwargs):
        self.socks_info = socks_info
        super().__init__(*args, **kwargs)


class SocksProxyHandler(BaseRequestHandler):
    def __init__(self, request_handler_class, socks_server_kwargs, *args, **kwargs):
        self.socks_kwargs = socks_server_kwargs or {}
        self.request_handler_class = request_handler_class
        super().__init__(*args, **kwargs)


class Socks5ProxyHandler(StreamRequestHandler, SocksProxyHandler):

    # SOCKS5 protocol https://tools.ietf.org/html/rfc1928
    # SOCKS5 username/password authentication https://tools.ietf.org/html/rfc1929

    def handle(self):
        sleep = self.socks_kwargs.get('sleep')
        if sleep:
            time.sleep(sleep)
        version, nmethods = self.connection.recv(2)
        assert version == SOCKS5_VERSION
        methods = list(self.connection.recv(nmethods))

        auth = self.socks_kwargs.get('auth')

        if auth is not None and Socks5Auth.AUTH_USER_PASS not in methods:
            self.connection.sendall(struct.pack('!BB', SOCKS5_VERSION, Socks5Auth.AUTH_NO_ACCEPTABLE))
            self.server.close_request(self.request)
            return

        elif Socks5Auth.AUTH_USER_PASS in methods:
            self.connection.sendall(struct.pack("!BB", SOCKS5_VERSION, Socks5Auth.AUTH_USER_PASS))

            _, user_len = struct.unpack('!BB', self.connection.recv(2))
            username = self.connection.recv(user_len).decode()
            pass_len = ord(self.connection.recv(1))
            password = self.connection.recv(pass_len).decode()

            if username == auth[0] and password == auth[1]:
                self.connection.sendall(struct.pack('!BB', SOCKS5_USER_AUTH_VERSION, SOCKS5_USER_AUTH_SUCCESS))
            else:
                self.connection.sendall(struct.pack('!BB', SOCKS5_USER_AUTH_VERSION, SOCKS5_USER_AUTH_FAILURE))
                self.server.close_request(self.request)
                return

        elif Socks5Auth.AUTH_NONE in methods:
            self.connection.sendall(struct.pack('!BB', SOCKS5_VERSION, Socks5Auth.AUTH_NONE))
        else:
            self.connection.sendall(struct.pack('!BB', SOCKS5_VERSION, Socks5Auth.AUTH_NO_ACCEPTABLE))
            self.server.close_request(self.request)
            return

        version, command, _, address_type = struct.unpack('!BBBB', self.connection.recv(4))
        socks_info = {
            'version': version,
            'auth_methods': methods,
            'command': command,
            'client_address': self.client_address,
            'ipv4_address': None,
            'domain_address': None,
            'ipv6_address': None,
        }
        if address_type == Socks5AddressType.ATYP_IPV4:
            socks_info['ipv4_address'] = socket.inet_ntoa(self.connection.recv(4))
        elif address_type == Socks5AddressType.ATYP_DOMAINNAME:
            socks_info['domain_address'] = self.connection.recv(ord(self.connection.recv(1))).decode()
        elif address_type == Socks5AddressType.ATYP_IPV6:
            socks_info['ipv6_address'] = socket.inet_ntop(socket.AF_INET6, self.connection.recv(16))
        else:
            self.server.close_request(self.request)

        socks_info['port'] = struct.unpack('!H', self.connection.recv(2))[0]

        # dummy response, the returned IP is just a placeholder
        self.connection.sendall(struct.pack(
            '!BBBBIH', SOCKS5_VERSION, self.socks_kwargs.get('reply', Socks5Reply.SUCCEEDED), 0x0, 0x1, 0x7f000001, 40000))

        self.request_handler_class(self.request, self.client_address, self.server, socks_info=socks_info)


class Socks4ProxyHandler(StreamRequestHandler, SocksProxyHandler):

    # SOCKS4 protocol http://www.openssh.com/txt/socks4.protocol
    # SOCKS4A protocol http://www.openssh.com/txt/socks4a.protocol

    def _read_until_null(self):
        return b''.join(iter(functools.partial(self.connection.recv, 1), b'\x00'))

    def handle(self):
        sleep = self.socks_kwargs.get('sleep')
        if sleep:
            time.sleep(sleep)
        socks_info = {
            'version': SOCKS4_VERSION,
            'command': None,
            'client_address': self.client_address,
            'ipv4_address': None,
            'port': None,
            'domain_address': None,
        }
        version, command, dest_port, dest_ip = struct.unpack('!BBHI', self.connection.recv(8))
        socks_info['port'] = dest_port
        socks_info['command'] = command
        if version != SOCKS4_VERSION:
            self.server.close_request(self.request)
            return
        use_remote_dns = False
        if 0x0 < dest_ip <= 0xFF:
            use_remote_dns = True
        else:
            socks_info['ipv4_address'] = socket.inet_ntoa(struct.pack("!I", dest_ip))

        user_id = self._read_until_null().decode()
        if user_id != (self.socks_kwargs.get('user_id') or ''):
            self.connection.sendall(struct.pack(
                '!BBHI', SOCKS4_REPLY_VERSION, Socks4CD.REQUEST_REJECTED_DIFFERENT_USERID, 0x00, 0x00000000))
            self.server.close_request(self.request)
            return

        if use_remote_dns:
            socks_info['domain_address'] = self._read_until_null().decode()

        # dummy response, the returned IP is just a placeholder
        self.connection.sendall(
            struct.pack(
                '!BBHI', SOCKS4_REPLY_VERSION,
                self.socks_kwargs.get('cd_reply', Socks4CD.REQUEST_GRANTED), 40000, 0x7f000001))

        self.request_handler_class(self.request, self.client_address, self.server, socks_info=socks_info)


class IPv6ThreadingTCPServer(ThreadingTCPServer):
    address_family = socket.AF_INET6


class SocksHTTPTestRequestHandler(http.server.BaseHTTPRequestHandler, SocksTestRequestHandler):
    def do_GET(self):
        if self.path == '/socks_info':
            payload = json.dumps(self.socks_info.copy())
            self.send_response(200)
            self.send_header('Content-Type', 'application/json; charset=utf-8')
            self.send_header('Content-Length', str(len(payload)))
            self.end_headers()
            self.wfile.write(payload.encode())


class SocksWebSocketTestRequestHandler(SocksTestRequestHandler):
    def handle(self):
        import websockets.sync.server
        protocol = websockets.ServerProtocol()
        connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=0)
        connection.handshake()
        connection.send(json.dumps(self.socks_info))
        connection.close()


@contextlib.contextmanager
def socks_server(socks_server_class, request_handler, bind_ip=None, **socks_server_kwargs):
    server = server_thread = None
    try:
        bind_address = bind_ip or '127.0.0.1'
        server_type = ThreadingTCPServer if '.' in bind_address else IPv6ThreadingTCPServer
        server = server_type(
            (bind_address, 0), functools.partial(socks_server_class, request_handler, socks_server_kwargs))
        server_port = http_server_port(server)
        server_thread = threading.Thread(target=server.serve_forever)
        server_thread.daemon = True
        server_thread.start()
        if '.' not in bind_address:
            yield f'[{bind_address}]:{server_port}'
        else:
            yield f'{bind_address}:{server_port}'
    finally:
        server.shutdown()
        server.server_close()
        server_thread.join(2.0)


class SocksProxyTestContext(abc.ABC):
    REQUEST_HANDLER_CLASS = None

    def socks_server(self, server_class, *args, **kwargs):
        return socks_server(server_class, self.REQUEST_HANDLER_CLASS, *args, **kwargs)

    @abc.abstractmethod
    def socks_info_request(self, handler, target_domain=None, target_port=None, **req_kwargs) -> dict:
        """return a dict of socks_info"""


class HTTPSocksTestProxyContext(SocksProxyTestContext):
    REQUEST_HANDLER_CLASS = SocksHTTPTestRequestHandler

    def socks_info_request(self, handler, target_domain=None, target_port=None, **req_kwargs):
        request = Request(f'http://{target_domain or "127.0.0.1"}:{target_port or "40000"}/socks_info', **req_kwargs)
        handler.validate(request)
        return json.loads(handler.send(request).read().decode())


class WebSocketSocksTestProxyContext(SocksProxyTestContext):
    REQUEST_HANDLER_CLASS = SocksWebSocketTestRequestHandler

    def socks_info_request(self, handler, target_domain=None, target_port=None, **req_kwargs):
        request = Request(f'ws://{target_domain or "127.0.0.1"}:{target_port or "40000"}', **req_kwargs)
        handler.validate(request)
        ws = handler.send(request)
        ws.send('socks_info')
        socks_info = ws.recv()
        ws.close()
        return json.loads(socks_info)


CTX_MAP = {
    'http': HTTPSocksTestProxyContext,
    'ws': WebSocketSocksTestProxyContext,
}


@pytest.fixture(scope='module')
def ctx(request):
    return CTX_MAP[request.param]()


class TestSocks4Proxy:
    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks4_no_auth(self, handler, ctx):
        with handler() as rh:
            with ctx.socks_server(Socks4ProxyHandler) as server_address:
                response = ctx.socks_info_request(
                    rh, proxies={'all': f'socks4://{server_address}'})
                assert response['version'] == 4

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks4_auth(self, handler, ctx):
        with handler() as rh:
            with ctx.socks_server(Socks4ProxyHandler, user_id='user') as server_address:
                with pytest.raises(ProxyError):
                    ctx.socks_info_request(rh, proxies={'all': f'socks4://{server_address}'})
                response = ctx.socks_info_request(
                    rh, proxies={'all': f'socks4://user:@{server_address}'})
                assert response['version'] == 4

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks4a_ipv4_target(self, handler, ctx):
        with ctx.socks_server(Socks4ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks4a://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='127.0.0.1')
                assert response['version'] == 4
                assert (response['ipv4_address'] == '127.0.0.1') != (response['domain_address'] == '127.0.0.1')

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks4a_domain_target(self, handler, ctx):
        with ctx.socks_server(Socks4ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks4a://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='localhost')
                assert response['version'] == 4
                assert response['ipv4_address'] is None
                assert response['domain_address'] == 'localhost'

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_ipv4_client_source_address(self, handler, ctx):
        with ctx.socks_server(Socks4ProxyHandler) as server_address:
            source_address = f'127.0.0.{random.randint(5, 255)}'
            with handler(proxies={'all': f'socks4://{server_address}'},
                         source_address=source_address) as rh:
                response = ctx.socks_info_request(rh)
                assert response['client_address'][0] == source_address
                assert response['version'] == 4

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    @pytest.mark.parametrize('reply_code', [
        Socks4CD.REQUEST_REJECTED_OR_FAILED,
        Socks4CD.REQUEST_REJECTED_CANNOT_CONNECT_TO_IDENTD,
        Socks4CD.REQUEST_REJECTED_DIFFERENT_USERID,
    ])
    def test_socks4_errors(self, handler, ctx, reply_code):
        with ctx.socks_server(Socks4ProxyHandler, cd_reply=reply_code) as server_address:
            with handler(proxies={'all': f'socks4://{server_address}'}) as rh:
                with pytest.raises(ProxyError):
                    ctx.socks_info_request(rh)

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_ipv6_socks4_proxy(self, handler, ctx):
        with ctx.socks_server(Socks4ProxyHandler, bind_ip='::1') as server_address:
            with handler(proxies={'all': f'socks4://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='127.0.0.1')
                assert response['client_address'][0] == '::1'
                assert response['ipv4_address'] == '127.0.0.1'
                assert response['version'] == 4

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_timeout(self, handler, ctx):
        with ctx.socks_server(Socks4ProxyHandler, sleep=2) as server_address:
            with handler(proxies={'all': f'socks4://{server_address}'}, timeout=0.5) as rh:
                with pytest.raises(TransportError):
                    ctx.socks_info_request(rh)


class TestSocks5Proxy:

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks5_no_auth(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh)
                assert response['auth_methods'] == [0x0]
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks5_user_pass(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler, auth=('test', 'testpass')) as server_address:
            with handler() as rh:
                with pytest.raises(ProxyError):
                    ctx.socks_info_request(rh, proxies={'all': f'socks5://{server_address}'})

                response = ctx.socks_info_request(
                    rh, proxies={'all': f'socks5://test:testpass@{server_address}'})

                assert response['auth_methods'] == [Socks5Auth.AUTH_NONE, Socks5Auth.AUTH_USER_PASS]
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks5_ipv4_target(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='127.0.0.1')
                assert response['ipv4_address'] == '127.0.0.1'
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks5_domain_target(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='localhost')
                assert (response['ipv4_address'] == '127.0.0.1') != (response['ipv6_address'] == '::1')
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks5h_domain_target(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks5h://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='localhost')
                assert response['ipv4_address'] is None
                assert response['domain_address'] == 'localhost'
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks5h_ip_target(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks5h://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='127.0.0.1')
                assert response['ipv4_address'] == '127.0.0.1'
                assert response['domain_address'] is None
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_socks5_ipv6_destination(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler) as server_address:
            with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='[::1]')
                assert response['ipv6_address'] == '::1'
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_ipv6_socks5_proxy(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler, bind_ip='::1') as server_address:
            with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
                response = ctx.socks_info_request(rh, target_domain='127.0.0.1')
                assert response['client_address'][0] == '::1'
                assert response['ipv4_address'] == '127.0.0.1'
                assert response['version'] == 5

    # XXX: is there any feasible way of testing IPv6 source addresses?
    # Same would go for non-proxy source_address test...
    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_ipv4_client_source_address(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler) as server_address:
            source_address = f'127.0.0.{random.randint(5, 255)}'
            with handler(proxies={'all': f'socks5://{server_address}'}, source_address=source_address) as rh:
                response = ctx.socks_info_request(rh)
                assert response['client_address'][0] == source_address
                assert response['version'] == 5

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Requests', 'http'), ('Websockets', 'ws')], indirect=True)
    @pytest.mark.parametrize('reply_code', [
        Socks5Reply.GENERAL_FAILURE,
        Socks5Reply.CONNECTION_NOT_ALLOWED,
        Socks5Reply.NETWORK_UNREACHABLE,
        Socks5Reply.HOST_UNREACHABLE,
        Socks5Reply.CONNECTION_REFUSED,
        Socks5Reply.TTL_EXPIRED,
        Socks5Reply.COMMAND_NOT_SUPPORTED,
        Socks5Reply.ADDRESS_TYPE_NOT_SUPPORTED,
    ])
    def test_socks5_errors(self, handler, ctx, reply_code):
        with ctx.socks_server(Socks5ProxyHandler, reply=reply_code) as server_address:
            with handler(proxies={'all': f'socks5://{server_address}'}) as rh:
                with pytest.raises(ProxyError):
                    ctx.socks_info_request(rh)

    @pytest.mark.parametrize('handler,ctx', [('Urllib', 'http'), ('Websockets', 'ws')], indirect=True)
    def test_timeout(self, handler, ctx):
        with ctx.socks_server(Socks5ProxyHandler, sleep=2) as server_address:
            with handler(proxies={'all': f'socks5://{server_address}'}, timeout=1) as rh:
                with pytest.raises(TransportError):
                    ctx.socks_info_request(rh)


if __name__ == '__main__':
    unittest.main()