"""Exercises the relay the way two phones would: one hosts, one joins by code, then
each sends game lines and expects the other to receive them untouched.

Start the relay first, then run this:

    python server/relay.py --host 127.0.0.1 &
    python server/test_relay.py
"""

import os
import shutil
import socket
import ssl
import subprocess
import sys
import tempfile
import time

HOST = "127.0.0.1"
PORT = 8642


def connect(port=PORT):
    s = socket.create_connection((HOST, port), timeout=5)
    return s, s.makefile("rwb")


def line(f):
    raw = f.readline()
    return raw.decode().strip() if raw else None


def send(f, text):
    f.write((text + "\n").encode())
    f.flush()


def check(name, got, expected):
    ok = got == expected
    print("%-42s %s" % (name, "ok" if ok else "FAILED  got %r want %r" % (got, expected)))
    return ok


failures = 0

# --- host opens a room ---
hs, hf = connect()
send(hf, "ROOM|HOST")
reply = line(hf)
failures += 0 if reply and reply.startswith("ROOM|OK|") else 1
code = reply.split("|")[2]
print("%-42s ok  (%s)" % ("host opened a room", code))

# --- a wrong code is refused, and does not disturb the room ---
ws, wf = connect()
send(wf, "ROOM|JOIN|NOSUCHCODE")
failures += 0 if check("wrong code refused", line(wf), "ROOM|ERR|no game with that code") else 1
ws.close()

# --- guest joins ---
gs, gf = connect()
send(gf, "ROOM|JOIN|%s" % code)
failures += 0 if check("guest accepted", line(gf), "ROOM|OK|%s" % code) else 1
failures += 0 if check("guest told it is paired", line(gf), "ROOM|PAIRED") else 1
failures += 0 if check("host told it is paired", line(hf), "ROOM|PAIRED") else 1

# --- a third player cannot squeeze in ---
ts, tf = connect()
send(tf, "ROOM|JOIN|%s" % code)
failures += 0 if check("third player refused", line(tf), "ROOM|ERR|that game is full") else 1
ts.close()

# --- game traffic is forwarded verbatim, both ways ---
send(hf, "HELLO|1|Andi")
failures += 0 if check("host -> guest", line(gf), "HELLO|1|Andi") else 1
send(gf, "FIRE|3|7")
failures += 0 if check("guest -> host", line(hf), "FIRE|3|7") else 1
send(hf, "REPORT|3|7|SUNK|Carrier|1")
failures += 0 if check("a whole report survives", line(gf), "REPORT|3|7|SUNK|Carrier|1") else 1

# --- a rapid burst keeps its order ---
for i in range(20):
    send(hf, "FIRE|%d|%d" % (i // 10, i % 10))
ordered = all(line(gf) == "FIRE|%d|%d" % (i // 10, i % 10) for i in range(20))
failures += 0 if check("twenty lines arrive in order", ordered, True) else 1

# --- leaving tells the other side ---
# makefile() dups the handle, so the socket is not really shut until both are closed.
hf.close()
hs.close()
time.sleep(0.4)
failures += 0 if check("guest told the host left", line(gf), "ROOM|GONE") else 1
gs.close()

# --- rubbish is rejected ---
bs, bf = connect()
send(bf, "hello there")
failures += 0 if check("rubbish rejected", line(bf), "ROOM|ERR|expected a room request") else 1
bs.close()

# --- an open relay tolerates a key line rather than choking on it ---
os_, of = connect()
send(of, "AUTH|whatever")
failures += 0 if check("open relay accepts a key line", line(of), "AUTH|OK") else 1
send(of, "ROOM|HOST")
reply = line(of)
failures += 0 if check("and then hands out a room", bool(reply and reply.startswith("ROOM|OK|")), True) else 1
os_.close()

# --- a relay started with a key serves only players who know it ---
KEY = "vier-woerter-die-keiner-raet"
KEY_PORT = PORT + 1
workdir = tempfile.mkdtemp()
keyfile = os.path.join(workdir, "key")
with open(keyfile, "w", encoding="utf-8") as handle:
    handle.write(KEY + chr(10))

keyed = subprocess.Popen(
    [sys.executable, os.path.join(os.path.dirname(os.path.abspath(__file__)), "relay.py"),
     "--host", HOST, "--port", str(KEY_PORT), "--key-file", keyfile],
    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
)
for _ in range(50):                       # give it a moment to bind
    try:
        socket.create_connection((HOST, KEY_PORT), timeout=1).close()
        break
    except OSError:
        time.sleep(0.1)

try:
    ks, kf = connect(KEY_PORT)
    send(kf, "AUTH|" + KEY)
    failures += 0 if check("right key accepted", line(kf), "AUTH|OK") else 1
    send(kf, "ROOM|HOST")
    reply = line(kf)
    failures += 0 if check("keyed relay hands out a room", bool(reply and reply.startswith("ROOM|OK|")), True) else 1
    keyed_code = reply.split("|")[2]

    ns, nf = connect(KEY_PORT)
    send(nf, "ROOM|HOST")
    failures += 0 if check("no key at all refused", line(nf),
                           "AUTH|ERR|this relay needs an access key") else 1
    ns.close()

    started = time.time()
    bs2, bf2 = connect(KEY_PORT)
    send(bf2, "AUTH|nearly-the-right-key")
    failures += 0 if check("wrong key refused", line(bf2), "AUTH|ERR|wrong access key") else 1
    failures += 0 if check("and costs a second", time.time() - started >= 1.0, True) else 1
    bs2.close()

    # Two players who both know the key still play exactly as before.
    gs2, gf2 = connect(KEY_PORT)
    send(gf2, "AUTH|" + KEY)
    line(gf2)
    send(gf2, "ROOM|JOIN|%s" % keyed_code)
    line(gf2)
    failures += 0 if check("keyed guest paired", line(gf2), "ROOM|PAIRED") else 1
    line(kf)                              # the host is told the same
    send(kf, "FIRE|1|1")
    failures += 0 if check("keyed traffic forwarded", line(gf2), "FIRE|1|1") else 1
    gs2.close()
    ks.close()
finally:
    keyed.terminate()
    keyed.wait(timeout=5)
    shutil.rmtree(workdir, ignore_errors=True)

# --- with a certificate, the whole conversation is encrypted ---
if shutil.which("openssl") is None:
    print("%-42s %s" % ("TLS checks", "skipped, no openssl on PATH"))
else:
    tlsdir = tempfile.mkdtemp()
    cert = os.path.join(tlsdir, "cert.pem")
    tlskey = os.path.join(tlsdir, "key.pem")
    subprocess.run(
        ["openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes", "-days", "1",
         "-subj", "/CN=localhost", "-addext", "subjectAltName=DNS:localhost",
         "-keyout", tlskey, "-out", cert],
        check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    TLS_PORT = PORT + 2
    encrypted = subprocess.Popen(
        [sys.executable, os.path.join(os.path.dirname(os.path.abspath(__file__)), "relay.py"),
         "--host", HOST, "--port", str(TLS_PORT), "--tls-cert", cert, "--tls-key", tlskey],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    for _ in range(50):
        try:
            socket.create_connection((HOST, TLS_PORT), timeout=1).close()
            break
        except OSError:
            time.sleep(0.1)

    try:
        trust = ssl.create_default_context(cafile=cert)

        def tls_connect():
            raw = socket.create_connection((HOST, TLS_PORT), timeout=5)
            wrapped = trust.wrap_socket(raw, server_hostname="localhost")
            return wrapped, wrapped.makefile("rwb")

        ts1, tf1 = tls_connect()
        failures += 0 if check("TLS handshake with a trusted cert", True, True) else 1
        send(tf1, "ROOM|HOST")
        reply = line(tf1)
        failures += 0 if check("room opened over TLS", bool(reply and reply.startswith("ROOM|OK|")), True) else 1
        tls_code = reply.split("|")[2]

        ts2, tf2 = tls_connect()
        send(tf2, "ROOM|JOIN|%s" % tls_code)
        line(tf2)
        failures += 0 if check("guest paired over TLS", line(tf2), "ROOM|PAIRED") else 1
        line(tf1)
        send(tf1, "FIRE|9|9")
        failures += 0 if check("traffic forwarded over TLS", line(tf2), "FIRE|9|9") else 1
        ts1.close()
        ts2.close()

        # A cert nobody vouches for must be refused, or TLS would be decoration.
        refused = False
        try:
            blind = ssl.create_default_context()
            raw = socket.create_connection((HOST, TLS_PORT), timeout=5)
            blind.wrap_socket(raw, server_hostname="localhost").close()
        except ssl.SSLError:
            refused = True
        failures += 0 if check("untrusted cert refused", refused, True) else 1

        # A client that speaks plain text to a TLS port gets nowhere.
        plain_failed = False
        try:
            ps, pf = connect(TLS_PORT)
            send(pf, "ROOM|HOST")
            plain_failed = line(pf) is None
            ps.close()
        except OSError:
            plain_failed = True
        failures += 0 if check("plain client refused by TLS port", plain_failed, True) else 1
    finally:
        encrypted.terminate()
        encrypted.wait(timeout=5)
        shutil.rmtree(tlsdir, ignore_errors=True)

print()
print("FAILURES: %d" % failures)
sys.exit(1 if failures else 0)