#!/usr/bin/env python3
import asyncio
import ipaddress
import json
import os
import signal
import struct
import time
from datetime import datetime
from typing import Optional, Tuple

SOCKS_VERSION = 5
CMD_CONNECT = 1
ATYP_IPV4 = 1
ATYP_DOMAIN = 3
ATYP_IPV6 = 4

REP_SUCCEEDED = 0
REP_GENERAL_FAILURE = 1
REP_CONNECTION_NOT_ALLOWED = 2
REP_NETWORK_UNREACHABLE = 3
REP_HOST_UNREACHABLE = 4
REP_COMMAND_NOT_SUPPORTED = 7
REP_ADDRESS_TYPE_NOT_SUPPORTED = 8

YGG_NET = ipaddress.ip_network("200::/7")


def env_endpoint(var_name: str, default: str) -> Tuple[str, int]:
    raw = os.getenv(var_name, default).strip()
    host, sep, port_s = raw.rpartition(":")
    if not sep:
        raise ValueError(f"Invalid endpoint for {var_name}: {raw}")
    return host, int(port_s)


LISTEN_HOST = os.getenv("ROUTER_LISTEN_HOST", "0.0.0.0")
LISTEN_PORT = int(os.getenv("ROUTER_LISTEN_PORT", "1080"))
UPSTREAM_TOR = env_endpoint("ROUTER_UPSTREAM_TOR", "10.5.0.7:9050")
UPSTREAM_I2P = env_endpoint("ROUTER_UPSTREAM_I2P", "10.5.0.2:4447")
UPSTREAM_YGG = env_endpoint("ROUTER_UPSTREAM_YGG", "127.0.0.1:1085")
I2P_MAP_FILE = os.getenv("ROUTER_I2P_MAP_FILE", "").strip()
I2P_POOL_CIDR = os.getenv("ROUTER_I2P_POOL_CIDR", "172.31.0.0/16")
I2P_POOL = ipaddress.ip_network(I2P_POOL_CIDR)

_i2p_map_mtime = -1.0
_i2p_map_last_check = 0.0
_i2p_ip_to_host = {}

BLOCKLIST = [
    item.strip().lower()
    for item in os.getenv("ROUTER_BLOCKLIST", "vk.com,ok.ru,mail.ru,gosuslugi.ru").split(",")
    if item.strip()
]


def log(msg: str) -> None:
    ts = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{ts}] [socks-router] {msg}", flush=True)


def is_blocked_domain(hostname: str) -> bool:
    host = hostname.rstrip(".").lower()
    for blocked in BLOCKLIST:
        if host == blocked or host.endswith("." + blocked):
            return True
    return False


def _reload_i2p_map_if_needed() -> None:
    global _i2p_map_mtime, _i2p_map_last_check, _i2p_ip_to_host

    if not I2P_MAP_FILE:
        return

    now = time.monotonic()
    if now - _i2p_map_last_check < 1.0:
        return
    _i2p_map_last_check = now

    try:
        mtime = os.path.getmtime(I2P_MAP_FILE)
    except OSError:
        _i2p_ip_to_host = {}
        _i2p_map_mtime = -1.0
        return

    if mtime == _i2p_map_mtime:
        return

    try:
        with open(I2P_MAP_FILE, "r", encoding="utf-8") as fh:
            data = json.load(fh)
        mapping = data.get("ip_to_host", {}) if isinstance(data, dict) else {}
        if isinstance(mapping, dict):
            _i2p_ip_to_host = {str(k): str(v) for k, v in mapping.items()}
        else:
            _i2p_ip_to_host = {}
        _i2p_map_mtime = mtime
    except Exception:
        _i2p_ip_to_host = {}


def lookup_i2p_hostname_for_virtual_ip(ip_str: str) -> Optional[str]:
    if not I2P_MAP_FILE:
        return None

    try:
        ip = ipaddress.ip_address(ip_str)
    except ValueError:
        return None

    if ip.version != 4 or ip not in I2P_POOL:
        return None

    _reload_i2p_map_if_needed()
    host = _i2p_ip_to_host.get(str(ip))
    if host and host.endswith(".i2p"):
        return host
    return None


def pick_upstream(host: str, atyp: int) -> Optional[Tuple[str, int]]:
    if atyp == ATYP_DOMAIN:
        dhost = host.rstrip(".").lower()
        if is_blocked_domain(dhost):
            return None
        if dhost.endswith(".onion"):
            return UPSTREAM_TOR
        if dhost.endswith(".i2p"):
            return UPSTREAM_I2P
        if dhost.endswith(".ygg") or dhost.endswith(".meshname") or dhost.endswith(".meship"):
            return UPSTREAM_YGG
        return UPSTREAM_TOR

    if atyp == ATYP_IPV6:
        try:
            ip = ipaddress.ip_address(host)
            if ip in YGG_NET:
                return UPSTREAM_YGG
        except ValueError:
            return None
        return UPSTREAM_TOR

    if atyp == ATYP_IPV4:
        return UPSTREAM_TOR

    return None


async def read_exact(reader: asyncio.StreamReader, n: int) -> bytes:
    data = await reader.readexactly(n)
    if len(data) != n:
        raise asyncio.IncompleteReadError(partial=data, expected=n)
    return data


async def send_reply(writer: asyncio.StreamWriter, rep: int) -> None:
    writer.write(b"\x05" + bytes([rep]) + b"\x00\x01\x00\x00\x00\x00\x00\x00")
    await writer.drain()


async def pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    try:
        while True:
            chunk = await reader.read(65536)
            if not chunk:
                break
            writer.write(chunk)
            await writer.drain()
    except Exception:
        pass
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


async def handle_client(client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
    peer = client_writer.get_extra_info("peername")
    upstream_writer = None

    try:
        header = await read_exact(client_reader, 2)
        ver, n_methods = header[0], header[1]
        if ver != SOCKS_VERSION:
            client_writer.close()
            await client_writer.wait_closed()
            return

        _ = await read_exact(client_reader, n_methods)
        client_writer.write(b"\x05\x00")
        await client_writer.drain()

        req = await read_exact(client_reader, 4)
        ver, cmd, _rsv, atyp = req
        if ver != SOCKS_VERSION or cmd != CMD_CONNECT:
            await send_reply(client_writer, REP_COMMAND_NOT_SUPPORTED)
            return

        if atyp == ATYP_IPV4:
            raw_addr = await read_exact(client_reader, 4)
            dst_host = str(ipaddress.IPv4Address(raw_addr))
        elif atyp == ATYP_IPV6:
            raw_addr = await read_exact(client_reader, 16)
            dst_host = str(ipaddress.IPv6Address(raw_addr))
        elif atyp == ATYP_DOMAIN:
            dlen = (await read_exact(client_reader, 1))[0]
            raw_addr = await read_exact(client_reader, dlen)
            dst_host = raw_addr.decode("ascii", errors="ignore")
        else:
            await send_reply(client_writer, REP_ADDRESS_TYPE_NOT_SUPPORTED)
            return

        dst_port = struct.unpack("!H", await read_exact(client_reader, 2))[0]

        connect_atyp = atyp
        connect_host = dst_host
        connect_raw_addr = raw_addr

        if atyp == ATYP_IPV4:
            mapped_i2p = lookup_i2p_hostname_for_virtual_ip(dst_host)
            if mapped_i2p:
                connect_atyp = ATYP_DOMAIN
                connect_host = mapped_i2p
                connect_raw_addr = None

        upstream = pick_upstream(connect_host, connect_atyp)
        if upstream is None:
            await send_reply(client_writer, REP_CONNECTION_NOT_ALLOWED)
            log(f"deny peer={peer} dst={dst_host}:{dst_port}")
            return

        up_host, up_port = upstream
        if connect_host != dst_host:
            log(
                f"route peer={peer} dst={dst_host}:{dst_port} mapped={connect_host} via={up_host}:{up_port}"
            )
        else:
            log(f"route peer={peer} dst={dst_host}:{dst_port} via={up_host}:{up_port}")

        try:
            upstream_reader, upstream_writer = await asyncio.open_connection(up_host, up_port)
        except Exception:
            await send_reply(client_writer, REP_NETWORK_UNREACHABLE)
            return

        # SOCKS5 greeting to upstream.
        upstream_writer.write(b"\x05\x01\x00")
        await upstream_writer.drain()
        up_auth = await read_exact(upstream_reader, 2)
        if up_auth[0] != SOCKS_VERSION or up_auth[1] != 0x00:
            await send_reply(client_writer, REP_GENERAL_FAILURE)
            return

        connect_req = bytearray(b"\x05\x01\x00")
        connect_req.append(connect_atyp)
        if connect_atyp == ATYP_DOMAIN:
            try:
                host_bytes = connect_host.encode("idna")
            except UnicodeError:
                host_bytes = connect_host.encode("ascii", errors="ignore")
            if len(host_bytes) > 255:
                await send_reply(client_writer, REP_HOST_UNREACHABLE)
                return
            connect_req.append(len(host_bytes))
            connect_req.extend(host_bytes)
        else:
            connect_req.extend(connect_raw_addr)
        connect_req.extend(struct.pack("!H", dst_port))

        upstream_writer.write(connect_req)
        await upstream_writer.drain()

        up_resp = await read_exact(upstream_reader, 4)
        if up_resp[0] != SOCKS_VERSION:
            await send_reply(client_writer, REP_GENERAL_FAILURE)
            return

        up_rep = up_resp[1]
        up_atyp = up_resp[3]
        if up_atyp == ATYP_IPV4:
            _ = await read_exact(upstream_reader, 4)
        elif up_atyp == ATYP_IPV6:
            _ = await read_exact(upstream_reader, 16)
        elif up_atyp == ATYP_DOMAIN:
            ln = (await read_exact(upstream_reader, 1))[0]
            _ = await read_exact(upstream_reader, ln)
        else:
            await send_reply(client_writer, REP_GENERAL_FAILURE)
            return
        _ = await read_exact(upstream_reader, 2)

        if up_rep != REP_SUCCEEDED:
            await send_reply(client_writer, up_rep if up_rep <= REP_ADDRESS_TYPE_NOT_SUPPORTED else REP_GENERAL_FAILURE)
            return

        await send_reply(client_writer, REP_SUCCEEDED)

        await asyncio.gather(
            pipe(client_reader, upstream_writer),
            pipe(upstream_reader, client_writer),
        )

    except asyncio.IncompleteReadError:
        pass
    except Exception as exc:
        log(f"error peer={peer} err={exc}")
        try:
            await send_reply(client_writer, REP_GENERAL_FAILURE)
        except Exception:
            pass
    finally:
        try:
            client_writer.close()
            await client_writer.wait_closed()
        except Exception:
            pass


async def main() -> int:
    server = await asyncio.start_server(handle_client, LISTEN_HOST, LISTEN_PORT)
    sockets = ", ".join(str(s.getsockname()) for s in (server.sockets or []))
    log(f"listening on {sockets}")

    stop_event = asyncio.Event()

    def _stop() -> None:
        stop_event.set()

    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            loop.add_signal_handler(sig, _stop)
        except NotImplementedError:
            pass

    await stop_event.wait()
    server.close()
    await server.wait_closed()
    log("shutdown complete")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(asyncio.run(main()))
    except KeyboardInterrupt:
        raise SystemExit(0)