"""A relay for Sea Duel games played over the internet.

MIT licensed - see the LICENSE file beside this one. Copyright (c) 2017 Danny Burt
for the original work, (c) 2026 Andreas Grauf for this.

Two players cannot usually reach each other directly: home routers do not forward
incoming connections. So both of them connect out to this, and it pipes bytes between
them. It understands nothing about the game - once two players are paired, every line
one sends is handed to the other untouched.

    python relay.py            # listens on 0.0.0.0:8642, open to anyone
    python relay.py --port 9000 --host 127.0.0.1
    python relay.py --key-file /etc/battleship-relay/key
    python relay.py --tls-cert /path/fullchain.pem --tls-key /path/privkey.pem

Standard library only, so it will run on anything with Python 3.8 or newer: a spare
machine, a Raspberry Pi, or the cheapest VPS going.

With --key-file the relay serves only players who know the key in that file, which is
what a relay on a public address wants: without it, anyone who finds the port can open
rooms on it. The key is compared in constant time, and a wrong one costs a second before
the connection is dropped, so guessing it over the network gets nowhere.

With --tls-cert and --tls-key the whole conversation is encrypted, key line included. The
handshake happens in the connection's own thread, so a client that cannot complete one
costs nobody else anything.

The protocol spoken before pairing, and never after:

    client -> AUTH|<key>                first line; always allowed, required with --key-file
    server -> AUTH|OK                   accepted, or accepted-and-ignored on an open relay
    server -> AUTH|ERR|<reason>         wrong key or none offered; the connection then ends
    client -> ROOM|HOST                 ask for a new room
    server -> ROOM|OK|FRIGATE           the code to pass to the other player
    client -> ROOM|JOIN|FRIGATE         join an existing room
    server -> ROOM|OK|FRIGATE           accepted
    server -> ROOM|PAIRED               sent to both; the game protocol starts now
    server -> ROOM|ERR|<reason>         no such room, room full, or malformed
    server -> ROOM|GONE                 the other player disconnected
"""

import argparse
import hmac
import random
import socket
import ssl
import threading
import time

DEFAULT_PORT = 8642
ROOM_IDLE_SECONDS = 15 * 60
MAX_LINE = 4096
# What a wrong key costs. Small enough to go unnoticed once, large enough that working
# through a keyspace one connection at a time is hopeless.
WRONG_KEY_SECONDS = 1.0
# A handshake that never finishes must not hold a thread for the room timeout.
HANDSHAKE_SECONDS = 20

# Short, memorable and unambiguous when read aloud or typed.
WORDS = [
    "ANCHOR", "BEACON", "CANNON", "DINGHY", "ENSIGN", "FRIGATE", "GALLEY", "HARBOUR",
    "ISLAND", "JETTY", "KEEL", "LAGOON", "MARINER", "NARWHAL", "OCEAN", "PENNANT",
    "QUAY", "RUDDER", "SEXTANT", "TIDE", "URCHIN", "VOYAGE", "WHARF", "YACHT",
]

rooms = {}
lock = threading.Lock()


class Room:
    def __init__(self, code):
        self.code = code
        self.host = None
        self.guest = None
        self.created = time.time()

    def other(self, connection):
        return self.guest if connection is self.host else self.host

    @property
    def full(self):
        return self.host is not None and self.guest is not None


def send(connection, line):
    try:
        connection.sendall((line + "\n").encode("utf-8"))
    except OSError:
        pass


def new_code():
    """A word, or a word with a digit once the plain words are taken."""
    for word in random.sample(WORDS, len(WORDS)):
        if word not in rooms:
            return word
    for _ in range(500):
        candidate = random.choice(WORDS) + str(random.randint(2, 99))
        if candidate not in rooms:
            return candidate
    return None


def read_line(connection, buffer):
    """Reads one line, keeping whatever came after it for the caller."""
    while b"\n" not in buffer:
        if len(buffer) > MAX_LINE:
            return None, buffer
        chunk = connection.recv(4096)
        if not chunk:
            return None, buffer
        buffer += chunk
    line, _, rest = buffer.partition(b"\n")
    return line.decode("utf-8", "replace").strip(), rest


def authenticate(connection, line, key):
    """Checks the shared key. Returns True if the caller may go on to ask for a room."""
    parts = line.split("|", 1)
    if parts[0] != "AUTH":
        if key is None:
            return True                      # An open relay does not care.
        send(connection, "AUTH|ERR|this relay needs an access key")
        return False

    offered = parts[1] if len(parts) > 1 else ""
    if key is not None and not hmac.compare_digest(offered, key):
        time.sleep(WRONG_KEY_SECONDS)
        send(connection, "AUTH|ERR|wrong access key")
        return False

    send(connection, "AUTH|OK")
    return True


def pair_up(connection, line, buffer, address):
    """Handles the room protocol. Returns the room once two players are in it."""
    parts = line.split("|")
    if len(parts) < 2 or parts[0] != "ROOM":
        send(connection, "ROOM|ERR|expected a room request")
        return None, buffer

    action = parts[1]
    with lock:
        if action == "HOST":
            code = new_code()
            if code is None:
                send(connection, "ROOM|ERR|no rooms free")
                return None, buffer
            room = Room(code)
            room.host = connection
            rooms[code] = room
            send(connection, "ROOM|OK|%s" % code)
            print("[%s] room %s opened" % (address, code), flush=True)
            return room, buffer

        if action == "JOIN":
            code = parts[2].strip().upper() if len(parts) > 2 else ""
            room = rooms.get(code)
            if room is None:
                send(connection, "ROOM|ERR|no game with that code")
                return None, buffer
            if room.full:
                send(connection, "ROOM|ERR|that game is full")
                return None, buffer
            room.guest = connection
            send(connection, "ROOM|OK|%s" % code)
            send(room.host, "ROOM|PAIRED")
            send(room.guest, "ROOM|PAIRED")
            print("[%s] joined room %s; playing" % (address, code), flush=True)
            return room, buffer

    send(connection, "ROOM|ERR|unknown request")
    return None, buffer


def serve(connection, address, key=None, tls=None):
    if tls is not None:
        # Wrapped here rather than around the listening socket: a client that cannot
        # handshake then costs its own thread and nobody else's.
        connection.settimeout(HANDSHAKE_SECONDS)
        try:
            connection = tls.wrap_socket(connection, server_side=True)
        except (OSError, ssl.SSLError) as problem:
            print("[%s] no TLS handshake: %s" % (address, problem), flush=True)
            connection.close()
            return

    connection.settimeout(ROOM_IDLE_SECONDS)
    buffer = b""
    room = None
    try:
        line, buffer = read_line(connection, buffer)
        if line is None:
            return

        # The key line is optional on an open relay, so whether it was sent is decided by
        # what the first line turns out to be rather than by asking for it.
        if not authenticate(connection, line, key):
            print("[%s] refused: bad or missing key" % address, flush=True)
            return
        if line.split("|", 1)[0] == "AUTH":
            line, buffer = read_line(connection, buffer)
            if line is None:
                return

        room, buffer = pair_up(connection, line, buffer, address)
        if room is None:
            return

        # Anything already buffered belongs to the game; forward it as it arrives.
        while True:
            if buffer:
                chunk, buffer = buffer, b""
            else:
                chunk = connection.recv(4096)
                if not chunk:
                    break
            with lock:
                other = room.other(connection)
            if other is not None:
                try:
                    other.sendall(chunk)
                except OSError:
                    break
    except (OSError, socket.timeout):
        pass
    finally:
        if room is not None:
            with lock:
                other = room.other(connection)
                if room.host is connection:
                    room.host = None
                elif room.guest is connection:
                    room.guest = None
                if room.host is None and room.guest is None:
                    rooms.pop(room.code, None)
                    print("[%s] room %s closed" % (address, room.code), flush=True)
            if other is not None:
                # Tell them, then close: a half-empty room has nothing left to relay, and
                # the client reads the end of the stream as the other player leaving.
                send(other, "ROOM|GONE")
                other.close()
        connection.close()


def sweep():
    """Drops rooms nobody ever joined, so codes do not leak away."""
    while True:
        time.sleep(60)
        cutoff = time.time() - ROOM_IDLE_SECONDS
        with lock:
            stale = [
                code for code, room in rooms.items()
                if not room.full and room.created < cutoff
            ]
            for code in stale:
                room = rooms.pop(code)
                for side in (room.host, room.guest):
                    if side is not None:
                        send(side, "ROOM|ERR|room expired")
                        side.close()
        if stale:
            print("swept %d stale room(s)" % len(stale), flush=True)


def main():
    parser = argparse.ArgumentParser(description="Battleship relay")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
    parser.add_argument(
        "--key-file",
        help="a file holding the access key; without it the relay serves anyone",
    )
    parser.add_argument("--tls-cert", help="certificate chain in PEM; needs --tls-key")
    parser.add_argument("--tls-key", help="private key in PEM; needs --tls-cert")
    args = parser.parse_args()

    if bool(args.tls_cert) != bool(args.tls_key):
        raise SystemExit("--tls-cert and --tls-key go together")

    tls = None
    if args.tls_cert:
        tls = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        tls.minimum_version = ssl.TLSVersion.TLSv1_2
        tls.load_cert_chain(args.tls_cert, args.tls_key)

    key = None
    if args.key_file:
        with open(args.key_file, encoding="utf-8") as handle:
            key = handle.read().strip()
        if not key:
            raise SystemExit("%s is empty - a blank key would let anyone in" % args.key_file)

    threading.Thread(target=sweep, daemon=True).start()

    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind((args.host, args.port))
    listener.listen(16)
    print("Battleship relay listening on %s:%d%s%s" % (
        args.host, args.port,
        ", key required" if key else ", open to anyone",
        ", TLS" if tls else ", unencrypted"), flush=True)

    while True:
        connection, address = listener.accept()
        threading.Thread(
            target=serve, args=(connection, "%s:%d" % address, key, tls), daemon=True
        ).start()


if __name__ == "__main__":
    main()