#!/usr/bin/env python3
"""A read-only FSP (File Service Protocol) server, for ROM streaming to Swiss.

    sudo ./scripts/fsp-server.py --root /workspace --password alltheroms

Swiss streams games over FSP/UDP, not TCP -- see cube/patches/bba/tcpip.c.
This serves a directory read-only so Swiss can browse it, stat a disc image
and stream it while a game runs.  It also reports the achieved read rate,
which is the number we actually want to know.

Protocol taken from Swiss's own two clients, not from memory:
  * cube/swiss/source/devices/fsp/fsplib.c   (the menu: browse, stat, read)
  * cube/patches/bba/tcpip.c                 (in-game streaming reads)

Wire format (v2, all multi-byte fields BIG-endian):

    cmd(1) checksum(1) key(2) sequence(2) len(2) pos(4)  then data[len] extra[xlen]

  checksum: sum = total_size; sum += every byte (with the checksum byte
            zeroed); sum += sum >> 8; take the low 8 bits.
  password: the data field is "path\\npassword\\0" -- the '\\n' REPLACES the
            path's terminating NUL (fsplib buildfilename).
  dir entry: time(4) size(4) type(1) name(NUL-terminated), 4-byte aligned;
            type 0x00 END, 0x01 FILE, 0x02 DIR, 0x2A SKIP.
  stat reply: the same 9-byte time/size/type prefix.

⚠️ GET_FILE replies must be a MULTIPLE OF 32 BYTES unless they are the last
piece of the request: the in-game client does
    data_size = MIN(requested, OSRoundDown32B(reply_len))
so a non-32-aligned reply is silently truncated down and the transfer stalls.
"""

import argparse
import os
import socket
import struct
import sys
import time

CC_VERSION, CC_INFO = 0x10, 0x11
CC_ERR, CC_GET_DIR, CC_GET_FILE = 0x40, 0x41, 0x42
CC_BYE, CC_STAT = 0x4A, 0x4D
CC_GET_PRO = 0x47
# Directory protection bits (fsplib.h).  Swiss uses these to decide whether
# the device supports patches:
#   if ((pro & (MKDIR|ADD|DEL|GET)) == (MKDIR|ADD|DEL) || (pro & OWNER))
#       features |= FEAT_PATCHES;  emulable |= EMU_ETHERNET;
# so a plain read-only answer makes Swiss DISABLE patch support here.
PRO_OWNER, PRO_DEL, PRO_ADD, PRO_MKDIR = 0x01, 0x02, 0x04, 0x08
PRO_GET, PRO_README, PRO_LIST, PRO_RENAME = 0x10, 0x20, 0x40, 0x80
PRO_BYTES = 1
RDTYPE_END, RDTYPE_FILE, RDTYPE_DIR, RDTYPE_SKIP = 0x00, 0x01, 0x02, 0x2A
HSIZE, FSP_SPACE = 12, 1024
SERVER_KEY = 0x1234


def checksum(pkt):
    """pkt must already have byte 1 (the checksum) zeroed."""
    total = len(pkt) + sum(pkt)
    total += total >> 8
    return total & 0xFF


def pack(cmd, key, seq, pos, data=b"", extra=b""):
    body = struct.pack("!BBHHHI", cmd, 0, key, seq, len(data), pos) + data + extra
    return body[:1] + bytes([checksum(body)]) + body[2:]


def unpack(pkt):
    if len(pkt) < HSIZE:
        return None
    cmd, csum, key, seq, ln, pos = struct.unpack("!BBHHHI", pkt[:HSIZE])
    if len(pkt) < HSIZE + ln:
        return None
    data = pkt[HSIZE:HSIZE + ln]
    extra = pkt[HSIZE + ln:]
    zeroed = pkt[:1] + b"\0" + pkt[2:]
    return dict(cmd=cmd, key=key, seq=seq, pos=pos, data=data, extra=extra,
                csum_ok=(checksum(zeroed) == csum))


def split_path(data, password):
    """Return (path, password_ok).  fsplib sends "path\\npassword\\0"."""
    s = data.split(b"\0")[0].decode("latin-1", "replace")
    if "\n" in s:
        path, _, pw = s.partition("\n")
    else:
        path, pw = s, ""
    return path, (password is None or pw == password)


def safe_join(root, path):
    p = os.path.normpath(os.path.join(root, path.lstrip("/")))
    if not (p == root or p.startswith(root + os.sep)):
        return None
    return p


def dir_block(real):
    """One FSP directory block: entries 4-byte aligned, then an END marker."""
    out = bytearray()
    try:
        names = sorted(os.listdir(real))
    except OSError:
        names = []
    for nm in names:
        full = os.path.join(real, nm)
        try:
            st = os.stat(full)
        except OSError:
            continue
        isdir = os.path.isdir(full)
        ent = struct.pack("!IIB", int(st.st_mtime) & 0xFFFFFFFF,
                          st.st_size & 0xFFFFFFFF,
                          RDTYPE_DIR if isdir else RDTYPE_FILE)
        ent += nm.encode("latin-1", "replace") + b"\0"
        while len(ent) % 4:
            ent += b"\0"
        # Keep the whole listing inside one block: the client treats
        # "fewer than 9 bytes left in this block" as SKIP.
        if len(out) + len(ent) + 9 > FSP_SPACE:
            break
        out += ent
    out += struct.pack("!IIB", 0, 0, RDTYPE_END)
    while len(out) % 4:
        out += b"\0"
    return bytes(out)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=".", help="directory to serve")
    ap.add_argument("--password", default=None,
                    help="require this password (Swiss sends it with every request)")
    ap.add_argument("--bind", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=21)
    ap.add_argument("--quiet-reads", action="store_true",
                    help="only log the rate, not every read")
    ap.add_argument("--claim-owner", action="store_true",
                    help="report FSP_DIR_OWNER on GET_PRO, which makes Swiss "
                         "enable FEAT_PATCHES and EMU_ETHERNET on this device. "
                         "⚠️ This server is READ-ONLY -- it has no UP_LOAD or "
                         "MAKE_DIR -- so Swiss believing it can write here "
                         "will make any patch write fail.  Try without it "
                         "first; use it only if a boot demands patches.")
    a = ap.parse_args()
    root = os.path.realpath(a.root)

    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        s.bind((a.bind, a.port))
    except PermissionError:
        sys.exit(f"port {a.port} needs root; run with sudo")
    print(f"fsp-server on {a.bind}:{a.port}  root {root}")
    print(f"  password {'set' if a.password else 'NONE'}")
    print(f"  point Swiss's FSP device at this host and port\n")

    served = 0
    t_first = None
    t_report = time.time()
    reads = 0
    fh_path, fh = None, None

    while True:
        pkt, peer = s.recvfrom(FSP_SPACE + HSIZE + 64)
        r = unpack(pkt)
        if r is None:
            print(f"  short packet from {peer[0]} ({len(pkt)} B)")
            continue
        if not r["csum_ok"]:
            print(f"  ⚠️  bad checksum from {peer[0]}, cmd {r['cmd']:#04x}")

        cmd, seq, pos = r["cmd"], r["seq"], r["pos"]

        if cmd == CC_VERSION:
            s.sendto(pack(CC_VERSION, SERVER_KEY, seq, 0,
                          b"fsp-server (rebbarb)\0"), peer)
            print(f"[{peer[0]}] VERSION")
            continue
        if cmd == CC_BYE:
            s.sendto(pack(CC_BYE, SERVER_KEY, seq, 0), peer)
            print(f"[{peer[0]}] BYE")
            continue

        path, pw_ok = split_path(r["data"], a.password)
        real = safe_join(root, path)
        if not pw_ok:
            print(f"[{peer[0]}] {cmd:#04x} {path!r} -> WRONG PASSWORD")
            s.sendto(pack(CC_ERR, SERVER_KEY, seq, pos, b"bad password\0"), peer)
            continue
        if real is None:
            s.sendto(pack(CC_ERR, SERVER_KEY, seq, pos, b"bad path\0"), peer)
            continue

        if cmd == CC_STAT:
            if os.path.exists(real):
                st = os.stat(real)
                body = struct.pack("!IIB", int(st.st_mtime) & 0xFFFFFFFF,
                                   st.st_size & 0xFFFFFFFF,
                                   RDTYPE_DIR if os.path.isdir(real)
                                   else RDTYPE_FILE)
                print(f"[{peer[0]}] STAT {path!r} -> {st.st_size} B")
            else:
                body = struct.pack("!IIB", 0, 0, 0)   # type 0 == ENOENT
                print(f"[{peer[0]}] STAT {path!r} -> not found")
            s.sendto(pack(CC_STAT, SERVER_KEY, seq, pos, body), peer)
            continue

        if cmd == CC_GET_PRO:
            # Reply shape (fsplib fsp_getpro): pos MUST equal FSP_PRO_BYTES,
            # and the flags byte sits at buf[len] -- i.e. the first byte of
            # the EXTRA region, after the echoed name, not inside data.
            # A non-existent directory must be CC_ERR so fsplib reports
            # ENOENT: Swiss probes "swiss/patches", then "swiss", then "".
            if not os.path.isdir(real):
                s.sendto(pack(CC_ERR, SERVER_KEY, seq, pos, b"no such dir\0"), peer)
                print(f"[{peer[0]}] GET_PRO {path!r} -> ENOENT (expected; "
                      f"Swiss falls back)")
                continue
            pro = PRO_GET | PRO_LIST | (PRO_OWNER if a.claim_owner else 0)
            s.sendto(pack(CC_GET_PRO, SERVER_KEY, seq, PRO_BYTES,
                          path.encode("latin-1", "replace"), bytes([pro])), peer)
            print(f"[{peer[0]}] GET_PRO {path!r} -> {pro:#04x}"
                  + ("  (OWNER: Swiss will enable patches)" if a.claim_owner
                     else "  (read-only: Swiss disables patches here)"))
            continue

        if cmd == CC_GET_DIR:
            if not os.path.isdir(real):
                s.sendto(pack(CC_ERR, SERVER_KEY, seq, pos, b"not a dir\0"), peer)
                continue
            blk = dir_block(real)
            # The client reads blocks until one comes back shorter than the
            # first; so the whole listing, then an empty reply.
            out = blk if pos == 0 else b""
            s.sendto(pack(CC_GET_DIR, SERVER_KEY, seq, pos, out), peer)
            print(f"[{peer[0]}] GET_DIR {path!r} pos={pos} -> {len(out)} B")
            continue

        if cmd == CC_GET_FILE:
            want = struct.unpack("!H", r["extra"][:2])[0] if len(r["extra"]) >= 2 else 1024
            want = min(want, FSP_SPACE)
            if fh_path != real:
                if fh:
                    fh.close()
                try:
                    fh = open(real, "rb")
                    fh_path = real
                except OSError as e:
                    s.sendto(pack(CC_ERR, SERVER_KEY, seq, pos,
                                  f"{e}\0".encode()), peer)
                    continue
            fh.seek(pos)
            data = fh.read(want)
            s.sendto(pack(CC_GET_FILE, SERVER_KEY, seq, pos, data), peer)
            served += len(data)
            reads += 1
            now = time.time()
            if t_first is None:
                t_first = now
                print(f"[{peer[0]}] GET_FILE {path!r} streaming started")
            if not a.quiet_reads and reads <= 3:
                print(f"[{peer[0]}] GET_FILE {path!r} pos={pos} "
                      f"want={want} -> {len(data)} B")
            if now - t_report >= 2.0:
                el = max(now - t_first, 1e-9)
                print(f"  streaming {served/el/1024:7.0f} KB/s   "
                      f"{served/2**20:8.1f} MiB in {el:6.1f}s   "
                      f"{reads/el:5.0f} reads/s   last pos {pos:#x}")
                t_report = now
            continue

        print(f"[{peer[0]}] unsupported command {cmd:#04x} {path!r}")
        s.sendto(pack(CC_ERR, SERVER_KEY, seq, pos, b"unsupported\0"), peer)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nstopped")