#!/usr/bin/env python3
import hashlib
import os
import signal
import socket
import struct
import subprocess
import sys
import tempfile
import threading
import time

failed = 0
procs = []


def check(name, cond):
    global failed
    print(("ok   " if cond else "FAIL ") + name)
    failed += not cond


def family(host):
    return socket.AF_INET6 if ":" in host else socket.AF_INET


def free_port(host="127.0.0.1"):
    s = socket.socket(family(host))
    s.bind((host, 0))
    port = s.getsockname()[1]
    s.close()
    return port


def start(tmp, listen, *args, lines=1):
    log = os.path.join(tmp, listen.replace("/", "_") + ".log")
    with open(log, "w") as f:
        p = subprocess.Popen(["./usocks", "-l", listen, *args], stderr=f)
    procs.append(p)
    for _ in range(100):
        with open(log) as f:
            if f.read().count("listening on") >= lines:
                return p
        if p.poll() is not None:
            break
        time.sleep(0.05)
    sys.exit(f"usocks -l {listen} did not start, see {log}")


def echo_server(host="127.0.0.1"):
    ls = socket.socket(family(host))
    ls.bind((host, 0))
    ls.listen(64)

    def serve(c):
        while b := c.recv(65536):
            c.sendall(b)
        c.shutdown(socket.SHUT_WR)
        c.close()

    def loop():
        while True:
            c, _ = ls.accept()
            threading.Thread(target=serve, args=(c,), daemon=True).start()

    threading.Thread(target=loop, daemon=True).start()
    return ls.getsockname()[1]


def recvn(s, n):
    b = b""
    while len(b) < n:
        x = s.recv(n - len(b))
        if not x:
            raise EOFError
        b += x
    return b


def recvall(s):
    b = b""
    while x := s.recv(65536):
        b += x
    return b


def request(host, port, cmd=1):
    for af, atyp in ((socket.AF_INET, 1), (socket.AF_INET6, 4)):
        try:
            a = bytes([atyp]) + socket.inet_pton(af, host)
            break
        except OSError:
            pass
    else:
        a = bytes([3, len(host)]) + host.encode()
    return bytes([5, cmd, 0]) + a + struct.pack(">H", port)


def reply(s):
    h = recvn(s, 4)
    recvn(s, {1: 4, 4: 16}[h[3]] + 2)
    return h[1]


def proxy(addr):
    if isinstance(addr, str):
        s = socket.socket(socket.AF_UNIX)
        s.connect(addr)
    elif isinstance(addr, tuple):
        s = socket.create_connection(addr)
    else:
        s = socket.create_connection(("127.0.0.1", addr))
    s.settimeout(30)
    return s


def connect(addr, host, port, auth=None, extra=b"", cmd=1):
    """Send the whole handshake in one write; return socket and reply code."""
    s = proxy(addr)
    if auth:
        u, p = auth
        msg = b"\5\1\2" + bytes([1, len(u)]) + u + bytes([len(p)]) + p
    else:
        msg = b"\5\1\0"
    s.sendall(msg + request(host, port, cmd) + extra)
    if recvn(s, 2) != (b"\5\2" if auth else b"\5\0"):
        return s, -1
    if auth and recvn(s, 2) != b"\1\0":
        return s, -2
    return s, reply(s)


def nfds(pid):
    return len(os.listdir(f"/proc/{pid}/fd"))


def proc(pid):
    try:
        with open(f"/proc/{pid}/stat") as f:
            s = f.read().rsplit(")", 1)[1].split()
    except OSError:
        return None
    return s[0], int(s[1])


def alive(pid):
    p = proc(pid)
    return p is not None and p[0] != "Z"


def workers(pid):
    ps = [int(d) for d in os.listdir("/proc") if d.isdigit()]
    return [k for k in ps if alive(k) and proc(k)[1] == pid]


def main():
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    echo = echo_server()
    tmp = tempfile.mkdtemp()
    unix = os.path.join(tmp, "usocks.sock")
    port, aport, mport, wport = (free_port() for _ in range(4))
    plain = start(tmp, f"127.0.0.1:{port}")
    start(tmp, f"127.0.0.1:{aport}", "-u", "bob", "-p", "secret")
    start(tmp, f"127.0.0.1:{mport}", "-m", "2")
    start(tmp, unix)
    fds = nfds(plain.pid)

    s, r = connect(port, "127.0.0.1", echo)
    s.sendall(b"hello")
    check("connect ipv4", r == 0 and recvn(s, 5) == b"hello")
    s.close()

    try:
        echo6, port6 = echo_server("::1"), free_port("::1")
    except OSError:
        print("skip ipv6: ::1 not available")
    else:
        start(tmp, f"[::1]:{port6}")
        s, r = connect(("::1", port6), "::1", echo6)
        s.sendall(b"six")
        check("ipv6 listener and destination", r == 0 and recvn(s, 3) == b"six")
        s.close()

    s, r = connect(port, "localhost", echo, extra=b"early")
    check("domain name with early data", r == 0 and recvn(s, 5) == b"early")
    s.close()

    s, r = connect(port, "nonexistent.invalid", echo)
    check("unresolvable name gets 4", r == 4)
    s, r = connect(port, "127.0.0.1", free_port())
    check("refused gets 5", r == 5)
    s, r = connect(port, "evil\nlog.example", echo)
    check("non-printable domain name gets 1", r == 1)
    s, r = connect(port, "127.0.0.1", echo, cmd=2)
    check("bind gets 7", r == 7)
    s = proxy(port)
    s.sendall(b"\5\1\0\5\1\0\7")
    recvn(s, 2)
    check("bad address type gets 8", reply(s) == 8)
    s = proxy(port)
    s.sendall(b"\4\1")
    check("bad version closes", recvall(s) == b"")
    s = proxy(port)
    s.sendall(b"\5\1\2")
    check("no acceptable method", recvn(s, 2) == b"\5\377")

    s = proxy(aport)
    for b in b"\5\1\2\1\3bob\6secret" + request("127.0.0.1", echo) + b"x":
        s.send(bytes([b]))
        time.sleep(0.002)
    ok = recvn(s, 4) == b"\5\2\1\0" and reply(s) == 0
    check("byte by byte handshake with auth", ok and recvn(s, 1) == b"x")
    s.close()
    s, r = connect(aport, "127.0.0.1", echo, auth=(b"bob", b"secret"))
    check("auth accepted", r == 0)
    s, r = connect(aport, "127.0.0.1", echo, auth=(b"bob", b"secre"))
    check("wrong password rejected", r == -2)
    s, r = connect(aport, "127.0.0.1", echo)
    check("no auth offered to auth server", r == -1)

    s, r = connect(port, "127.0.0.1", echo)
    s.sendall(b"abc")
    s.shutdown(socket.SHUT_WR)
    check("client half-close", recvall(s) == b"abc")
    ls = socket.socket()
    ls.bind(("127.0.0.1", 0))
    ls.listen(1)
    s, r = connect(port, "127.0.0.1", ls.getsockname()[1])
    a, _ = ls.accept()
    a.sendall(b"bye")
    a.shutdown(socket.SHUT_WR)
    got = recvall(s)
    s.sendall(b"still open")
    s.shutdown(socket.SHUT_WR)
    check("remote half-close", got == b"bye" and recvall(a) == b"still open")

    data = os.urandom(1 << 20) * 32
    s, r = connect(port, "127.0.0.1", echo)

    def write():
        s.sendall(data)
        s.shutdown(socket.SHUT_WR)

    threading.Thread(target=write).start()
    got = recvall(s)
    check(
        "32 MiB both ways",
        hashlib.sha256(got).digest() == hashlib.sha256(data).digest(),
    )

    s, r = connect(unix, "127.0.0.1", echo)
    s.sendall(b"unix")
    check("unix socket listener", r == 0 and recvn(s, 4) == b"unix")
    s.close()

    stale = os.path.join(tmp, "stale.sock")
    s = socket.socket(socket.AF_UNIX)
    s.bind(stale)
    s.close()
    start(tmp, stale)
    r = subprocess.run(
        ["./usocks", "-l", stale], stderr=subprocess.DEVNULL, timeout=5, check=False
    )
    s, rc = connect(stale, "127.0.0.1", echo)
    check("stale unix socket replaced, live one kept", r.returncode == 1 and rc == 0)
    s.close()

    a, _ = connect(mport, "127.0.0.1", echo)
    b, _ = connect(mport, "127.0.0.1", echo)
    x = proxy(mport)
    x.sendall(b"\5\1\0")
    x.settimeout(0.5)
    try:
        x.recv(2)
        waited = False
    except TimeoutError:
        waited = True
    a.close()
    x.settimeout(5)
    check("maxconn queues the third client", waited and recvn(x, 2) == b"\5\0")

    w = start(tmp, f"127.0.0.1:{wport}", "-w", "2", lines=2)
    kids = workers(w.pid)
    base = [nfds(k) for k in kids]
    held = [connect(wport, "127.0.0.1", echo) for _ in range(16)]
    used = [nfds(k) > n for k, n in zip(kids, base)]
    ok = all(r == 0 for _, r in held) and len(kids) == 2 and all(used)
    check("workers share the port", ok)
    os.kill(kids[0], signal.SIGKILL)
    rc = w.wait(5)
    t = time.time()
    while alive(kids[1]) and time.time() - t < 5:
        time.sleep(0.05)
    check("a dead worker stops all", rc == 1 and not alive(kids[1]))

    # A listener with a full accept queue drops SYNs: connect never finishes.
    full = socket.socket()
    full.bind(("127.0.0.1", 0))
    full.listen(0)
    fill = []
    for _ in range(4):
        f = socket.socket()
        f.setblocking(False)
        f.connect_ex(full.getsockname())
        fill.append(f)
    silent = proxy(port)
    t = time.time()
    s, r = connect(port, "127.0.0.1", full.getsockname()[1])
    check("connect timeout gets 6", r == 6 and 9 < time.time() - t < 13)
    check("handshake timeout closes", recvall(silent) == b"")

    time.sleep(0.5)
    check("no leaked descriptors", nfds(plain.pid) == fds)
    print("FAILED" if failed else "all ok")
    sys.exit(1 if failed else 0)


try:
    main()
finally:
    for p in procs:
        p.terminate()
        p.wait()