"""Close behaviour of a 3proxy binary on Windows, for comparing two builds.

Usage: python close_test.py path\\to\\3proxy.exe label

1. upload (#638): a client sends a file through an HTTP CONNECT tunnel and
   closes as soon as the last byte is handed to its socket; the server
   counts what arrives.
2. download: the server sends a file and closes; the client counts it.
3. reply while uploading: the server reads part of a large upload, answers
   and closes; the client counts the answer while still sending.
4. handles: the proxy's handle count before and after many connections
   which close with data still queued.
5. timing: many short sequential connections through the proxy.
"""
import socket, subprocess, sys, threading, time, os, tempfile

BIN, LABEL = sys.argv[1], sys.argv[2]
FILE = os.urandom(1580872)


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


def start_proxy():
    port = free_port()
    d = tempfile.mkdtemp()
    cfg = os.path.join(d, "3proxy.cfg")
    with open(cfg, "w") as f:
        f.write("auth iponly\nallow *\nproxy -p%d\n" % port)
    p = subprocess.Popen([BIN, cfg])
    for _ in range(100):
        try:
            socket.create_connection(("127.0.0.1", port), 1).close()
            break
        except OSError:
            time.sleep(0.05)
    time.sleep(0.5)
    return p, port


def tunnel(pport, dport):
    c = socket.create_connection(("127.0.0.1", pport), 10)
    c.sendall(("CONNECT 127.0.0.1:%d HTTP/1.1\r\nHost: 127.0.0.1:%d\r\n\r\n" % (dport, dport)).encode())
    head = b""
    while b"\r\n\r\n" not in head:
        d = c.recv(1)
        if not d:
            raise OSError("tunnel refused")
        head += d
    if b" 200" not in head.split(b"\r\n")[0]:
        raise OSError("tunnel refused: %r" % head[:40])
    return c


def listener():
    l = socket.socket()
    l.bind(("127.0.0.1", 0))
    l.listen(50)
    return l, l.getsockname()[1]


def handles(pid):
    try:
        out = subprocess.check_output(["powershell", "-NoProfile", "-Command",
                                        "(Get-Process -Id %d).HandleCount" % pid])
        return int(out.strip())
    except Exception:
        return -1


def upload(pport, n=20):
    full = 0
    sizes = []
    for _ in range(n):
        l, dport = listener()
        got = []

        def srv():
            c, _ = l.accept()
            c.settimeout(15)
            total = 0
            try:
                while True:
                    d = c.recv(65536)
                    if not d:
                        break
                    total += len(d)
            except OSError:
                pass
            got.append(total)
            c.close()

        t = threading.Thread(target=srv)
        t.start()
        try:
            c = tunnel(pport, dport)
            c.sendall(FILE)
            c.close()
        except OSError:
            pass
        t.join(20)
        l.close()
        size = got[0] if got else -1
        sizes.append(size)
        full += size == len(FILE)
    return "upload %d/%d complete, short ones: %s" % (
        full, n, [s for s in sizes if s != len(FILE)][:6])


def download(pport, n=20):
    full = 0
    sizes = []
    for _ in range(n):
        l, dport = listener()

        def srv():
            c, _ = l.accept()
            try:
                c.sendall(FILE)
            except OSError:
                pass
            c.close()

        t = threading.Thread(target=srv)
        t.start()
        total = 0
        try:
            c = tunnel(pport, dport)
            c.settimeout(15)
            while True:
                d = c.recv(65536)
                if not d:
                    break
                total += len(d)
            c.close()
        except OSError:
            pass
        t.join(20)
        l.close()
        sizes.append(total)
        full += total == len(FILE)
    return "download %d/%d complete, short ones: %s" % (
        full, n, [s for s in sizes if s != len(FILE)][:6])


def reply_while_uploading(pport, n=10):
    full = 0
    errs = set()
    for _ in range(n):
        l, dport = listener()
        reply = b"R" * 65536

        def srv():
            c, _ = l.accept()
            got = 0
            while got < 16384:
                d = c.recv(16384 - got)
                if not d:
                    break
                got += len(d)
            c.sendall(reply)
            c.shutdown(socket.SHUT_WR)
            c.settimeout(10)
            try:
                while c.recv(65536):
                    pass
            except OSError:
                pass
            c.close()

        t = threading.Thread(target=srv)
        t.start()
        rx = 0
        try:
            c = tunnel(pport, dport)

            def send():
                try:
                    c.sendall(b"x" * 1500000)
                    c.shutdown(socket.SHUT_WR)
                except OSError as e:
                    errs.add("send:" + type(e).__name__)

            st = threading.Thread(target=send)
            st.start()
            c.settimeout(15)
            try:
                while True:
                    d = c.recv(65536)
                    if not d:
                        break
                    rx += len(d)
            except OSError as e:
                errs.add("recv:" + type(e).__name__)
            st.join(20)
            c.close()
        except OSError as e:
            errs.add("tunnel:" + type(e).__name__)
        t.join(20)
        l.close()
        full += rx == len(reply)
    return "reply while uploading %d/%d complete, errors %s" % (full, n, sorted(errs) or "none")


def handle_growth(p, pport, n=300):
    l, dport = listener()
    stop = []

    def srv():
        l.settimeout(1)
        while not stop:
            try:
                c, _ = l.accept()
            except OSError:
                continue
            # read a little and close with the rest queued and unread
            try:
                c.recv(1024)
            except OSError:
                pass
            c.close()

    t = threading.Thread(target=srv)
    t.start()
    before = handles(p.pid)
    for _ in range(n):
        try:
            c = tunnel(pport, dport)
            c.setblocking(False)
            try:
                c.send(b"y" * 262144)
            except OSError:
                pass
            c.close()
        except OSError:
            pass
    time.sleep(8)
    after = handles(p.pid)
    stop.append(1)
    t.join(5)
    l.close()
    return "handles before %d, after %d connections %d" % (before, n, after)


def timing(pport, n=300):
    l, dport = listener()

    def srv():
        for _ in range(n):
            c, _ = l.accept()
            d = c.recv(16)
            c.sendall(d)
            c.close()

    t = threading.Thread(target=srv)
    t.start()
    t0 = time.time()
    for _ in range(n):
        c = tunnel(pport, dport)
        c.sendall(b"ping")
        c.recv(16)
        c.close()
    dt = time.time() - t0
    t.join(10)
    l.close()
    return "timing %d sequential tunnels: %.2fs (%.1f ms each)" % (n, dt, dt * 1000 / n)


p, pport = start_proxy()
try:
    for f in (upload, download, reply_while_uploading, timing):
        print("[%s] %s" % (LABEL, f(pport)), flush=True)
    print("[%s] %s" % (LABEL, handle_growth(p, pport)), flush=True)
finally:
    p.kill()