#!/usr/bin/env python3
"""Online update client for RustROS.

  rustros-update.py HOST pair [--key-file F]           pair with a machine whose console opened pairing
                                                     (`pair` there): both show a check code; after `pair yes`
                                                     on the console the new key is saved to the key file
  rustros-update.py HOST status
  rustros-update.py HOST upload rustros.elf [--wait]   write into the spare slot, reboot into it once
  rustros-update.py HOST confirm                      keep the running (trial) kernel
  rustros-update.py HOST reboot
  rustros-update.py HOST poweroff
  rustros-update.py HOST module-config NAME [KEY [VALUE | -]]   show / set / remove a module setting
  rustros-update.py HOST net ip [nicN] add|del a.b.c.d[-e][,...]   extra addresses; also "net ip ...", "net mtu ...", "net save"
  rustros-update.py HOST mod NAME ARGS...                a module's own command (e.g. "mod ctnet ranges")
  rustros-update.py HOST deploy [BUILD] [--programs a,b] [--no-kernel]
                                                     a release onto the machine with one reboot: every module
                                                     it has (from BUILD/modules: one built for another kernel
                                                     interface is kept aside for that kernel), the programs,
                                                     then the kernel (BUILD/rustros.efi), its trial boot
                                                     confirmed; then checks no module was refused;
                                                     rescue, dhcp and the updater always (they replace the
                                                     boot image's while it runs); --no-kernel: no reboot
  rustros-update.py HOST net boot [KEY VALUE|-]          network settings for the next boot only (grubenv net*)
  rustros-update.py HOST log [prev|this|BOOT [LINES]]     console logs of the last boots (kept on disk)
  rustros-update.py HOST info                         uptime, slot, address, storage, modules, last crash
  rustros-update.py HOST module-install NAME FILE     install a module: native (.mod) or WebAssembly (.wasm);
                                                     replaces a running one, no reboot
  rustros-update.py HOST module-remove NAME [native|wasm]
  rustros-update.py HOST module-switch NAME native|wasm
  rustros-update.py HOST run PROGRAM [ARGS...] [-t] [--net NAME]
                                                     start a program file (static ELF; an absolute path that is
                                                     no file here: one on the machine) as a process; -t: on a
                                                     terminal of its own, interactively (Ctrl-] hangs up)
  rustros-update.py HOST docker ARGS...               docker on the machine (/usr/lib/rustros/docker), on a
                                                     terminal: interactive when this one is (docker exec -it)
  rustros-update.py HOST program-install NAME FILE    a program shipped with RustROS into /usr/lib/rustros/NAME
  rustros-update.py HOST esp-get FILE                 copy the EFI partition (the loaders, the slots, the boot state) to FILE
  rustros-update.py HOST esp-put FILE                 write back the sectors of FILE that differ (then verify)
  rustros-update.py HOST migrate KERNEL.efi [LOADER.efi] [--size-mb N] [--bios]
                                                     GRUB's layout to the boot plan's (--bios: also the
                                                     RustROS BIOS loader into the MBR): a new
                                                     EFI partition with the RustROS loader, the PE kernel and
                                                     the boot state (GRUB's settings), then one GPT change and
                                                     a restart; GRUB's partition stays where it was

The key comes from --key, $RUSTROS_KEY or --key-file (default build/update.key).
`upload --wait` waits for the trial boot to come up and confirms it.
"""

import argparse
import hashlib
import hmac
import os
import socket
import struct
import sys
import time


def recv_exact(s, n):
    buf = b""
    while len(buf) < n:
        chunk = s.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("connection closed")
        buf += chunk
    return buf


def hello(host, port, timeout=10):
    s = socket.create_connection((host, port), timeout=timeout)
    h = recv_exact(s, 22)
    if h[:4] != b"RRS1":
        raise ConnectionError("not a RustROS update server")
    return s, h[4:20], chr(h[20]), h[21] == 1


def keystream_xor(key, nonce, data):
    out = bytearray(data)
    for i in range(0, len(out), 32):
        k = hmac.new(key, nonce + b"E" + struct.pack(">I", i // 32), hashlib.sha256).digest()
        for j in range(min(32, len(out) - i)):
            out[i + j] ^= k[j]
    return bytes(out)


def program(path, args, net="", on_machine=False):
    """'G' / 'T' payload: the arguments (argv[0] the file's name), the ELF;
    `on_machine`: no ELF, argv[0] is the path of a program on the machine."""
    args = [path if on_machine else os.path.basename(path)] + list(args)
    if net:
        args = [f"--net={net}"] + args
    return bytes([len(args)]) + b"".join(bytes([len(x.encode())]) + x.encode() for x in args) + (b"" if on_machine else open(path, "rb").read())


class Term:
    """A program on a terminal of the machine ('T'): typed bytes go in,
    what it writes comes out, in encrypted and authenticated frames."""

    def __init__(self, host, port, key, path, args=(), rows=24, cols=80, net="", on_machine=False):
        s, nonce, _, _ = hello(host, port)
        payload = struct.pack(">HH", rows, cols) + program(path, args, net, on_machine)
        head = b"T" + struct.pack(">I", len(payload))
        s.sendall(head + hmac.new(key, nonce + b"H" + head, hashlib.sha256).digest() + payload + hmac.new(key, nonce + head + payload, hashlib.sha256).digest())
        line = b""
        while not line.endswith(b"\n"):
            c = s.recv(1)
            if not c:
                break
            line += c
        if not line.startswith(b"OK "):
            s.close()
            raise ConnectionError(line.decode(errors="replace").strip() or "connection closed")
        self.pid = int(line.split()[1])
        self.s, self.key, self.nonce = s, key, nonce
        self.rx = self.tx = 0
        self.buf = b""
        self.status = None  # the program's wait status once it ended

    def _xor(self, d, ctr, data):
        out = bytearray(data)
        for i in range(0, len(out), 32):
            k = hmac.new(self.key, self.nonce + b"T" + d + struct.pack(">QI", ctr, i // 32), hashlib.sha256).digest()
            for j in range(min(32, len(out) - i)):
                out[i + j] ^= k[j]
        return bytes(out)

    def _mac(self, d, ctr, frame):
        return hmac.new(self.key, self.nonce + b"t" + d + struct.pack(">Q", ctr) + frame, hashlib.sha256).digest()

    def send(self, kind, data=b""):
        for k in range(0, max(len(data), 1), 16384):
            f = kind + struct.pack(">H", len(data[k:k + 16384])) + self._xor(b"c", self.tx, data[k:k + 16384])
            self.s.sendall(f + self._mac(b"c", self.tx, f))
            self.tx += 1

    def type(self, data):
        self.send(b"i", data.encode() if isinstance(data, str) else data)

    def resize(self, rows, cols):
        self.send(b"w", struct.pack(">HH", rows, cols))

    def hangup(self):
        self.send(b"h")

    def read(self, timeout=1.0):
        """What the program wrote (b"" if nothing came in time); None once it ended."""
        out = b""
        self.s.settimeout(timeout)
        while True:
            while len(self.buf) >= 3:
                n = struct.unpack(">H", self.buf[1:3])[0]
                if len(self.buf) < 3 + n + 32:
                    break
                f, tag, self.buf = self.buf[:3 + n], self.buf[3 + n:3 + n + 32], self.buf[3 + n + 32:]
                if not hmac.compare_digest(tag, self._mac(b"s", self.rx, f)):
                    raise ConnectionError("bad frame MAC")
                data = self._xor(b"s", self.rx, f[3:])
                self.rx += 1
                if f[:1] == b"o":
                    out += data
                elif f[:1] == b"x":
                    self.status = struct.unpack(">i", data)[0]
            if out or self.status is not None:
                return out if out or self.status is None else None
            try:
                c = self.s.recv(65536)
            except socket.timeout:
                return out
            if not c:
                return out or None
            self.buf += c

    def expect(self, what, timeout=10.0):
        """Read until `what` (bytes or str) came: (found, everything read)."""
        what = what.encode() if isinstance(what, str) else what
        got, end = b"", time.time() + timeout
        while time.time() < end:
            d = self.read(min(0.5, max(0.01, end - time.time())))
            if d is None:
                break
            got += d
            if what in got:
                return True, got
        return what in got, got


def collect(t):
    """Without a local terminal: what `t` writes, until it ends."""
    while t.status is None:
        d = t.read(1.0)
        if d is None:
            break
        os.write(sys.stdout.fileno(), d)
    return t.status


def interactive(t):
    """The local terminal attached to `t` until the program ends (Ctrl-] hangs up)."""
    import select, signal, termios, tty
    fd = sys.stdin.fileno()
    old = termios.tcgetattr(fd)
    size = lambda: os.get_terminal_size(fd)
    signal.signal(signal.SIGWINCH, lambda *_: t.resize(size().lines, size().columns))
    tty.setraw(fd)
    try:
        while t.status is None:
            r, _, _ = select.select([fd, t.s], [], [], 1.0)
            if fd in r:
                data = os.read(fd, 4096)
                if b"\x1d" in data:
                    t.hangup()
                    break
                t.type(data)
            if t.s in r or t.buf:
                d = t.read(0.01)
                if d is None:
                    break
                os.write(sys.stdout.fileno(), d)
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old)
    return t.status


def command(host, port, key, cmd, payload=b"", encrypt=False):
    s, nonce, slot, trial = hello(host, port)
    command.nonce = nonce
    if encrypt:
        # 'E': the inner command travels encrypted, the MAC covers the ciphertext
        payload = keystream_xor(key, nonce, cmd + payload)
        cmd = b"E"
    head = cmd + struct.pack(">I", len(payload))
    hdr_mac = hmac.new(key, nonce + b"H" + head, hashlib.sha256).digest()
    mac = hmac.new(key, nonce + head + payload, hashlib.sha256).digest()
    s.sendall(head + hdr_mac)
    sent = 0
    t0 = time.time()
    while sent < len(payload):
        n = s.send(payload[sent:sent + 65536])
        sent += n
        if len(payload) > 65536:
            print(f"\r  {sent * 100 // len(payload)}%  {sent // 1024} KiB", end="", flush=True)
    if len(payload) > 65536:
        print(f"  ({time.time() - t0:.1f}s)")
    s.sendall(mac)
    # on a slow link most of the payload may still be on its way: the answer
    # comes once it has arrived and been written
    s.settimeout(max(s.gettimeout() or 10, 30 + len(payload) // 4096))
    reply = b""
    try:
        # 'I' and listings send several lines and close; the others answer one line
        while cmd in (b"I", b"E", b"N", b"L", b"J") or not reply.endswith(b"\n"):
            chunk = s.recv(256)
            if not chunk:
                break
            reply += chunk
    finally:
        s.close()
    return reply.decode(errors="replace").strip(), slot, trial


# The EFI partition: where the GPT says (sector 64 in today's images, 576 in
# GRUB's, elsewhere after a migration; kernels without boot-area reads had it
# at 576). Its length is in its FAT boot sector.
ESP_FIRST = 576
ESP_CHUNK = 128
ESP_TYPE = bytes.fromhex("28732ac11ff8d211ba4b00a0c93ec93b")
BOOT_AREA_END = 64 * 2048


def sector_read(a, key, lba, n, inner=b"D"):
    """Sectors of the EFI partition ('D') or of the boot area ('d'); None on an error reply."""
    reply, _, _ = command(a.host, a.port, key, inner, struct.pack("<QH", lba, n), encrypt=True)
    if not reply.startswith("OK "):
        return None
    return keystream_xor(key, command.nonce + b"R", bytes.fromhex(reply[3:]))


def esp_read(a, key, lba, n):
    data = sector_read(a, key, lba, n)
    if data is None:
        sys.exit(f"esp read at {lba} failed")
    return data


def gpt_parts(a, key):
    """The partitions (type, first, last, index) from the primary GPT, None if the
    kernel cannot read the boot area."""
    head = sector_read(a, key, 1, 33, b"d")
    if head is None or head[:8] != b"EFI PART":
        return None
    entries_lba, count, size = struct.unpack_from("<QII", head, 72)
    raw = head[(entries_lba - 1) * 512:] if entries_lba <= 33 else b""
    parts = []
    for i in range(min(count, len(raw) // size)):
        e = raw[i * size:(i + 1) * size]
        if any(e[:16]):
            first, last = struct.unpack_from("<QQ", e, 32)
            parts.append((e[:16], first, last, i))
    return parts


def esp_first(a, key):
    parts = gpt_parts(a, key)
    esp = [p for p in parts or [] if p[0] == ESP_TYPE]
    return esp[0][1] if esp else ESP_FIRST


def esp_get(a, key, first=None):
    first = esp_first(a, key) if first is None else first
    boot = esp_read(a, key, first, 1)
    total = struct.unpack_from("<H", boot, 19)[0] or struct.unpack_from("<I", boot, 32)[0]
    if boot[510:512] != b"\x55\xaa" or not total:
        sys.exit("esp: no FAT boot sector where the EFI partition should start")
    last = first + total - 1
    data = b""
    for lba in range(first, last + 1, ESP_CHUNK):
        data += esp_read(a, key, lba, min(ESP_CHUNK, last + 1 - lba))
    return data


def esp(a, key):
    first = esp_first(a, key)
    if a.action == "esp-get":
        data = esp_get(a, key, first)
        fd = os.open(a.image, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)  # it holds the key
        os.write(fd, data)
        os.close(fd)
        print(f"{len(data)} bytes ({len(data) // 512} sectors from {first})")
        return 0
    new = open(a.image, "rb").read()
    old = esp_get(a, key, first)
    if len(new) != len(old):
        sys.exit(f"{a.image} is {len(new)} bytes, the EFI partition {len(old)}")
    changed = [i for i in range(len(old) // 512) if old[i * 512:(i + 1) * 512] != new[i * 512:(i + 1) * 512]]
    # runs of changed sectors, at most ESP_CHUNK each
    runs = []
    for i in changed:
        if runs and runs[-1][1] == i and runs[-1][1] - runs[-1][0] < ESP_CHUNK:
            runs[-1][1] = i + 1
        else:
            runs.append([i, i + 1])
    for f, end in runs:
        reply, _, _ = command(a.host, a.port, key, b"B", struct.pack("<Q", first + f) + new[f * 512:end * 512], encrypt=True)
        if reply != "OK":
            sys.exit(f"esp write at {first + f}: {reply}")
    if esp_get(a, key, first) != new:
        sys.exit("verify failed: the EFI partition does not match the file")
    print(f"{len(changed)} sectors written in {len(runs)} runs, verified")
    return 0


class EspFat:
    """A FAT12/16 EFI partition (its bytes), for reading its files (migrate
    reads GRUB's partition with it)."""

    def __init__(self, img):
        self.img = bytearray(img)
        bps, self.spc, rsvd, self.nfats, root_entries = struct.unpack_from("<HBHBH", img, 11)
        self.fatsz = struct.unpack_from("<H", img, 22)[0]
        total = struct.unpack_from("<H", img, 19)[0] or struct.unpack_from("<I", img, 32)[0]
        if bps != 512 or not self.fatsz:
            sys.exit("not a FAT12/16 partition with 512-byte sectors")
        self.fat_start = rsvd
        self.root = rsvd + self.nfats * self.fatsz
        self.data = self.root + (root_entries * 32 + 511) // 512
        self.clusters = (total - self.data) // self.spc
        self.bits = 12 if self.clusters < 4085 else 16
        self.eoc = 0xFFF if self.bits == 12 else 0xFFFF

    def get(self, c):
        o = self.fat_start * 512
        if self.bits == 12:
            v = struct.unpack_from("<H", self.img, o + c * 3 // 2)[0]
            return v >> 4 if c & 1 else v & 0xFFF
        return struct.unpack_from("<H", self.img, o + c * 2)[0]

    def chain(self, c):
        out = []
        while 2 <= c < self.eoc - 7 and len(out) <= self.clusters:
            out.append(c)
            c = self.get(c)
        return out

    def lba(self, c):
        """The partition-relative sector of cluster `c`."""
        return self.data + (c - 2) * self.spc

    def entries(self, start):
        """(name, attr, first cluster, size, byte offset of the entry) of a
        directory (None: the root directory), without . and .. or long names."""
        if start is None:
            spans = [(self.root * 512, self.data * 512)]
        else:
            spans = [(self.lba(c) * 512, (self.lba(c) + self.spc) * 512) for c in self.chain(start)]
        for lo, hi in spans:
            for o in range(lo, hi, 32):
                e = self.img[o:o + 32]
                if e[0] == 0:
                    return
                if e[0] == 0xE5 or e[11] == 0x0F or e[:2] in (b". ", b".."):
                    continue
                first = struct.unpack_from("<H", e, 20)[0] << 16 | struct.unpack_from("<H", e, 26)[0]
                yield bytes(e[:11]), e[11], first, struct.unpack_from("<I", e, 28)[0], o

    def find(self, path):
        d, found = None, None
        for p in path.strip("/").split("/"):
            base, _, ext = p.upper().partition(".")
            want = (base.ljust(8) + ext.ljust(3)).encode()
            found = next((e for e in self.entries(d) if e[0] == want), None)
            if found is None:
                sys.exit(f"{path} not in the EFI partition")
            d = found[2]
        return found

    def read(self, path):
        _, _, first, size, _ = self.find(path)
        return b"".join(bytes(self.img[self.lba(c) * 512:(self.lba(c) + self.spc) * 512]) for c in self.chain(first))[:size]

def grub_settings(f):
    """From GRUB's EFI partition (an EspFat): the boot state's settings, the
    same ones grub.cfg passes today (grubenv's variables, the update key, the
    image's network defaults, the fixed kernel arguments)."""
    import re
    env = f.read("/boot/grub/grubenv").decode(errors="replace")
    cfg = f.read("/boot/grub/grub.cfg").decode(errors="replace")
    vars = {}
    for line in env.splitlines():
        k, eq, v = line.partition("=")
        if eq and not k.startswith("#") and k not in ("boot_slot", "try_slot", "boot_count", "rolled_back"):
            vars[k] = v
    # the image's defaults: if [ -z "$net_ip" ]; then set net_ip=...; fi
    for k, v in re.findall(r'if \[ -z "\$(net_\w+)" \]; then set \1=(\S+); fi', cfg):
        vars.setdefault(k, v)
    m = re.search(r"multiboot2 /boot/kernel-\$run_slot\.elf (.*)", cfg)
    if not m:
        sys.exit("grub.cfg has no RustROS menu entry")
    fixed = []
    for arg in m.group(1).split():
        if arg.startswith(("rustros.slot=", "rustros.trial=", "rustros.boots=", "rustros.slots=", "rustros.env=", "rustros.bundle=", "$")):
            continue  # the loader's, the slots' (in the boot state now), the bundle (installed), grubenv's
        if arg.startswith("key="):
            if not arg[4:].startswith("$"):
                vars.setdefault("update_key", arg[4:])
            continue
        fixed.append(arg)
    vars["cmdline"] = " ".join(fixed)
    if not vars.get("update_key"):
        sys.exit("no update key in grub.cfg or grubenv")
    return vars


def fat_files(f, d=None, path=""):
    """Every file of an EspFat: (path, bytes)."""
    for name, attr, first, size, _ in f.entries(d):
        if attr & 0x08:
            continue  # the volume label
        base, ext = name[:8].decode().rstrip(), name[8:].decode().rstrip()
        p = f"{path}/{base}" + (f".{ext}" if ext else "")
        if attr & 0x10:
            yield from fat_files(f, first, p)
        else:
            yield p, f.read(p)


def migrate(a, key, kernel_file, loader_file, mb):
    """GRUB's layout to the boot plan's, on a running UEFI machine: a new EFI
    partition (tools/fat.py: the RustROS loader as BOOTX64.EFI, the PE kernel
    in both slots, the boot state with GRUB's settings, the old partition's
    other files such as \\NvVars) is written into free space in the boot
    area and read back; then one GPT change points the EFI partition at it,
    and the machine restarts. The old partition stays where it was (the
    GPT moved back to it boots GRUB again). Nothing is changed before the
    GPT write: a power cut until then leaves GRUB's layout."""
    import fat, bootstate
    s, _, slot, trial = hello(a.host, a.port)
    s.close()
    if trial:
        sys.exit("confirm the running trial kernel first")
    kernel, loader = open(kernel_file, "rb").read(), open(loader_file, "rb").read()
    if kernel[:2] != b"MZ":
        sys.exit(f"{kernel_file}: the RustROS loader starts PE kernel files (build/rustros.efi)")
    parts = gpt_parts(a, key)
    if parts is None:
        sys.exit("this kernel cannot read the boot area; update it first")
    esp = [p for p in parts if p[0] == ESP_TYPE]
    if not esp:
        sys.exit("no EFI partition in the GPT")
    _, old_first, old_last, _ = esp[0]
    old = EspFat(esp_get(a, key, old_first))
    names = [p for p, _ in fat_files(old)]
    if "/boot/grub/grub.cfg" not in [n.lower() for n in names]:
        sys.exit("the EFI partition is not GRUB's (already migrated?)")
    vars = grub_settings(old)
    # the new place: after everything in the boot area, 1 MiB aligned
    sectors = mb * 2048
    end = max(p[2] for p in parts if p[1] < BOOT_AREA_END) + 1
    first = (end + 2047) // 2048 * 2048
    last = first + sectors - 1
    if last >= BOOT_AREA_END or any(p[1] <= last and first <= p[2] for p in parts):
        sys.exit(f"no room for {mb} MiB in the boot area (free from sector {end})")
    f = fat.Fat.format(sectors, hidden=first)
    for d in ("/EFI", "/EFI/BOOT", "/EFI/RUSTROS"):
        f.mkdir(d)
    f.write("/EFI/BOOT/BOOTX64.EFI", loader)
    ext = [f.write(f"/EFI/RUSTROS/KERNEL-{s}.EFI", kernel) for s in "AB"]
    # the old partition's other files: what the firmware keeps there (\NvVars)
    kept = []
    for p, data in fat_files(old):
        if p.upper().startswith(("/BOOT/", "/EFI/BOOT/")):
            continue  # GRUB, its configuration, the old kernel slots
        parent = p.rsplit("/", 1)[0]
        if parent:
            try:
                f.mkdir(parent)
            except AssertionError:
                pass  # there already
        f.write(p, data)
        kept.append(p)
    st = bootstate.encode(1, [(len(kernel), ext[0]), (len(kernel), ext[1])], vars)
    f.img[bootstate.AT[0]:bootstate.AT[0] + bootstate.SIZE] = st
    mbr = None
    if a.bios:
        # the RustROS BIOS loader: stage 2 twice into the new partition's
        # reserved sectors, stage 1 into the MBR's boot code (its partition
        # table and disk signature stay); GRUB's MBR is kept in a file
        import mkimg
        old_mbr = sector_read(a, key, 0, 1, b"d")
        if old_mbr is None or old_mbr[510:512] != b"\x55\xaa":
            sys.exit("cannot read the MBR (this kernel's storage service is too old?)")
        keep = f"mbr-backup-{a.host}.bin"
        open(keep, "wb").write(old_mbr)
        print(f"  GRUB's MBR saved in {keep}")
        mbr = mkimg.bios_loader(f.img, old_mbr, first)
    new = bytes(f.img)
    print(f"EFI partition: {(old_last - old_first + 1) // 2} KiB at {old_first} (GRUB) -> {mb} MiB at {first} (the RustROS loader)")
    print(f"  kernel {len(kernel)} bytes in both slots; kept from the old partition: {', '.join(kept) or 'nothing'}")
    print(f"  settings: {', '.join(sorted(k for k in vars if k != 'update_key'))}, update_key")
    step = ESP_CHUNK - 8
    writes = 0
    for i in range(0, sectors, step):
        if a.stop_after is not None and writes >= a.stop_after:
            sys.exit(f"stopped after {writes} writes (--stop-after)")
        chunk = new[i * 512:(i + step) * 512]
        reply, _, _ = command(a.host, a.port, key, b"b", struct.pack("<Q", first + i) + chunk, encrypt=True)
        if reply != "OK":
            sys.exit(f"write at {first + i}: {reply} (nothing changed: GRUB's partition is still in use)")
        writes += 1
    back = b"".join(sector_read(a, key, first + i, min(ESP_CHUNK, sectors - i), b"d") or b"" for i in range(0, sectors, ESP_CHUNK))
    if back != new:
        sys.exit("read back differs (nothing changed: GRUB's partition is still in use)")
    # grubenv may have changed since it was read: then the settings are stale
    if grub_settings(EspFat(esp_get(a, key, old_first))) != vars:
        sys.exit("the settings changed meanwhile: run it again (nothing changed)")
    if a.stop_after is not None and writes >= a.stop_after:
        sys.exit(f"stopped before the GPT write (--stop-after)")
    if mbr:
        # first the MBR (from now on the BIOS starts the RustROS loader, which
        # finds the new partition by its own sector numbers), then the GPT
        reply, _, _ = command(a.host, a.port, key, b"b", struct.pack("<Q", 0) + mbr, encrypt=True)
        if reply != "OK" or sector_read(a, key, 0, 1, b"d") != mbr:
            sys.exit(f"MBR write: {reply} (the MBR may be GRUB's still; run it again)")
        print("the MBR starts the RustROS BIOS loader now")
    reply, _, _ = command(a.host, a.port, key, b"M", struct.pack("<QQ", first, last), encrypt=True)
    if reply != "OK":
        sys.exit(f"move: {reply} (nothing changed)")
    print("the GPT now points at the new EFI partition; restarting into the RustROS loader")
    command(a.host, a.port, key, b"R")
    time.sleep(5)
    for _ in range(180):
        try:
            s, _, cur_slot, _ = hello(a.host, a.port, timeout=3)
            s.close()
            print(f"up again on slot {cur_slot}")
            return 0
        except OSError:
            time.sleep(1)
    sys.exit(f"it did not come back within 3 minutes (GRUB's partition is at {old_first}-{old_last}: "
             f"esp_move back to it from a rescue boot restores GRUB)")


def net_apply(a, key, secs):
    """As iptables-apply: show the changes, apply, then check that NEW
    connections still work before keeping them; otherwise the node reverts
    by itself when `secs` run out."""
    net = lambda text: command(a.host, a.port, key, b"N", text.encode())[0]
    preview = net("cand")
    print(preview)
    if not preview.startswith("candidate"):
        return 1
    if not a.yes and input("apply these changes? [y/N] ").strip().lower() not in ("y", "yes"):
        print("not applied (the candidate stays)")
        return 1
    t0 = time.time()
    r = net(f"apply {secs}")
    if r.startswith("ERR"):
        print(r)
        return 1
    print(f"applied; testing a new connection (reverted by itself in {secs} s unless kept) ...")
    try:
        s, _, _, _ = hello(a.host, a.port, timeout=max(1, min(5, secs - 1)) if secs else 5)
        s.close()
    except OSError as e:
        print(f"cannot open a new connection ({e}): the node goes back to the previous configuration")
        return 1
    if not a.yes:
        left = secs - (time.time() - t0) if secs else None
        q = "new connections work. keep these changes? [y/N]" + (f" ({left:.0f} s left) " if left else " ")
        if input(q).strip().lower() not in ("y", "yes"):
            print(net("revert"))
            return 1
    r = net("confirm")
    print(r)
    if r.startswith("too late") or r.startswith("nothing"):
        return 1
    print(f"keep it across reboots with: {sys.argv[0]} {a.host} net save")
    return 0


def inst(a, key):
    """System instances: inst ls | create k=v... /prog args... | start ID |
    exec ID [k=v...] /prog args... | signal ID SIG | kill ID | rm ID |
    budget ID k=v... (settings: name root id host net mem threads cpu cwd
    uid gid stdio env=...; the program is the first word starting with /)."""
    words = [w for w in [a.image, a.module, *a.value] if w]
    if not words:
        sys.exit("inst ls | create | start | exec | signal | kill | rm | budget")
    head, lines, prog = ["inst"], [], None
    for w in words:
        if prog is not None:
            prog.append(w)
        elif w.startswith("/"):
            prog = [w]
        elif "=" in w and len(head) > 1:
            lines.append(w)
        else:
            head.append(w)
    if prog:
        lines += [f"path={prog[0]}"] + [f"arg={x}" for x in prog]
    reply, _, _ = command(a.host, a.port, key, b"J", "\n".join([" ".join(head)] + lines).encode())
    print(reply, end="" if reply.endswith("\n") else "\n")
    return 1 if reply.startswith("ERR") else 0


def upload(a, key, path, wait):
    """A kernel file into the spare slot; with `wait`, its trial boot confirmed."""
    image = open(path, "rb").read()
    is_pe = image[:2] == b"MZ" and image[int.from_bytes(image[0x3C:0x40], "little"):][:4] == b"PE\0\0"
    if image[:4] != b"\x7fELF" and not is_pe:
        sys.exit(f"{path} is not a kernel file (ELF or PE)")
    print(f"uploading {path} ({len(image)} bytes)")
    reply, slot, _ = command(a.host, a.port, key, b"U", image)
    print(reply)
    if reply != "OK":
        return 1
    new_slot = "b" if slot == "a" else "a"
    if not wait:
        print(f"rebooting into slot {new_slot}; run '{sys.argv[0]} {a.host} confirm' once it is up")
        return 0
    print(f"waiting for slot {new_slot} to come up...")
    went_down = False  # the old kernel may still answer for a moment
    t0 = time.time()
    for _ in range(180):
        try:
            s, _, cur, trial = hello(a.host, a.port, timeout=3)
            s.close()
            if cur == new_slot and trial:
                break
            if cur == slot and not trial and (went_down or time.time() - t0 > 60):
                sys.exit(f"came back on the old slot {slot}: the new kernel failed to boot")
        except OSError:
            went_down = True
        time.sleep(1)
    else:
        sys.exit("timed out waiting for the new kernel")
    reply, _, _ = command(a.host, a.port, key, b"C")
    print(f"confirm: {reply}")
    return 0 if reply == "OK" else 1


def installed_modules(host, port, key):
    """The modules `info` lists: name -> status line."""
    reply, _, _ = command(host, port, key, b"I")
    mods = {}
    for l in reply.splitlines():
        if l.startswith("module   ") and ":" in l:
            name, rest = l[len("module   "):].split(":", 1)
            mods[name.strip()] = rest.strip()
    return mods


def deploy(a, key, build):
    """Modules first (the module manager keeps those built for another
    kernel interface aside), programs, then the kernel: one reboot, after
    which every module runs the version for the new kernel."""
    mods = installed_modules(a.host, a.port, key)
    # the boot image's apps always (their installed versions replace them
    # while the system runs); the updater last: it answers this tool
    builtins = [m for m in ("rescue", "dhcp", "updater") if os.path.exists(f"{build}/modules/{m}.mod")]
    first = ["storage", "vfs", "ext4", "sites"]
    order = [m for m in first if m in mods] + sorted(m for m in mods if m not in first and m not in ("netstack", "linux", *builtins)) \
        + [m for m in ("netstack", "linux") if m in mods] + builtins
    bad = []
    for m in order:
        f = f"{build}/modules/{m}.mod"
        if not os.path.exists(f):
            print(f"{m}: no {f}, left as it is")
            continue
        data = open(f, "rb").read()
        name = m.encode()
        for attempt in (1, 2):
            try:
                reply, _, _ = command(a.host, a.port, key, b"W", bytes([len(name)]) + name + data)
                break
            except OSError as e:
                reply = f"ERR {e}"
        print(f"{m}: {reply}")
        if m in builtins and "is a built-in app" in reply:
            # a kernel from before installed versions replaced them: the new kernel carries them
            print(f"{m}: the running kernel takes no installed {m}; the new kernel brings its own")
            continue
        if reply != "OK":
            bad.append(m)
        if m == "updater" and reply == "OK":
            time.sleep(4)  # it takes over a moment after answering
    for p in [x for x in a.programs.split(",") if x]:
        data = open(f"{build}/userland/{p}", "rb").read()
        name = p.encode()
        reply, _, _ = command(a.host, a.port, key, b"Y", bytes([len(name)]) + name + data)
        print(f"program {p}: {reply}")
        if reply != "OK":
            bad.append(p)
    if bad:
        print(f"not installed: {', '.join(bad)}; the kernel is not uploaded")
        return 1
    if a.no_kernel:
        print("deployed: modules only, no reboot")
        return 0
    r = upload(a, key, f"{build}/rustros.efi", True)
    if r != 0:
        return r
    # every module the machine had runs again, none refused
    time.sleep(5)
    after = installed_modules(a.host, a.port, key)
    reply, _, _ = command(a.host, a.port, key, b"L", b"this")
    refused = [l for l in reply.splitlines() if "built for another" in l or "binary interface" in l]
    for l in refused:
        print(l)
    missing = [m for m in mods if m not in after]
    if missing:
        print(f"gone after the update: {', '.join(missing)}")
    stack = after.get("netstack", "")
    if "netstack" in mods and not stack.startswith("native service"):
        print(f"netstack: {stack} (the installed stack did not take over)")
    ok = not refused and not missing and ("netstack" not in mods or stack.startswith("native service"))
    print("deployed: one reboot, every module on the new kernel" if ok else "deployed, with the problems above")
    return 0 if ok else 1


# X25519 (RFC 7748), for pairing
P25519 = 2 ** 255 - 19


def x25519(k, u):
    k = bytearray(k)
    k[0] &= 248
    k[31] = (k[31] & 127) | 64
    k = int.from_bytes(k, "little")
    x1 = int.from_bytes(u, "little") & ((1 << 255) - 1)
    x2, z2, x3, z3, swap = 1, 0, x1, 1, 0
    for t in reversed(range(255)):
        bit = (k >> t) & 1
        swap ^= bit
        if swap:
            x2, x3, z2, z3 = x3, x2, z3, z2
        swap = bit
        a, b = (x2 + z2) % P25519, (x2 - z2) % P25519
        aa, bb = a * a % P25519, b * b % P25519
        e = (aa - bb) % P25519
        c, d = (x3 + z3) % P25519, (x3 - z3) % P25519
        da, cb = d * a % P25519, c * b % P25519
        x3, z3 = (da + cb) ** 2 % P25519, x1 * (da - cb) ** 2 % P25519
        x2, z2 = aa * bb % P25519, e * (aa + 121665 * e) % P25519
    if swap:
        x2, z2 = x3, z3
    return (x2 * pow(z2, P25519 - 2, P25519) % P25519).to_bytes(32, "little")


def pair(a):
    """Pair with the machine (its console opened the window): the new key goes to the key file."""
    secret = os.urandom(32)
    client = x25519(secret, (9).to_bytes(32, "little"))
    s, nonce, _, _ = hello(a.host, a.port, timeout=150)
    s.sendall(b"A" + struct.pack(">I", 32) + bytes(32) + client)
    head = recv_exact(s, 4)
    if head != b"PAIR":
        sys.exit((head + s.recv(200)).decode(errors="replace").strip())
    server = recv_exact(s, 32)
    shared = x25519(secret, server)
    if shared == bytes(32):
        sys.exit("pairing: bad public key from the machine")
    t = nonce + client + server
    key = hmac.new(shared, b"rustros pair key" + t, hashlib.sha256).digest()[:16].hex()
    code = int.from_bytes(hmac.new(shared, b"rustros pair code" + t, hashlib.sha256).digest()[:4], "big") % 1000000
    print(f"check code {code // 1000:03} {code % 1000:03}")
    print("if the machine's console shows the same code, type there:  pair yes", flush=True)
    reply = b""
    while not reply.endswith(b"\n"):
        c = s.recv(100)
        if not c:
            break
        reply += c
    s.close()
    reply = reply.decode(errors="replace").strip()
    if reply != "OK":
        sys.exit(f"pairing: {reply or 'the connection closed'}")
    if os.path.exists(a.key_file) and os.path.getsize(a.key_file):
        os.replace(a.key_file, a.key_file + ".old")
    fd = os.open(a.key_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    os.write(fd, (key + "\n").encode())
    os.close(fd)
    print(f"paired: the update key is in {a.key_file}")


def main():
    # docker's own arguments are everything after "docker"
    argv, dargs = sys.argv[1:], []
    k, valued = 1, ("--port", "--key", "--key-file", "--stop-after", "--size-mb", "--net")
    while k < len(argv) and argv[k].startswith("-"):
        k += 2 if argv[k] in valued else 1
    if argv[k:k + 1] == ["docker"]:
        argv, dargs = argv[:k + 1], argv[k + 1:]
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("host")
    ap.add_argument("action", choices=["status", "upload", "confirm", "reboot", "poweroff", "info", "net", "mod", "log",
                                       "module-install", "module-remove", "module-switch", "module-config", "run", "ps", "kill", "inst",
                                       "esp-get", "esp-put", "migrate", "docker", "program-install", "deploy", "pair"])
    ap.add_argument("image", nargs="?", help="kernel image, or the module name for module-*")
    ap.add_argument("module", nargs="?", help="module file (module-install), native|wasm, or the setting's key (module-config)")
    ap.add_argument("value", nargs="*", help="module-config: the value ('-' removes the setting)")
    ap.add_argument("--port", type=int, default=9000)
    ap.add_argument("--key")
    ap.add_argument("--key-file", default="build/update.key")
    ap.add_argument("--wait", action="store_true", help="after upload, wait for the new kernel and confirm it")
    ap.add_argument("--stop-after", type=int, help="migrate: stop after this many writes (power-cut tests)")
    ap.add_argument("--size-mb", type=int, default=1, help="migrate: the new EFI partition's size (default 1 MiB, as new images)")
    ap.add_argument("--bios", action="store_true", help="migrate: a BIOS machine (the RustROS BIOS loader into the MBR and the partition)")
    ap.add_argument("--net", default="", help="run: in this network domain (a container network's own stack; made if new)")
    ap.add_argument("-t", "--tty", action="store_true", help="run: on a terminal, interactively")
    ap.add_argument("-y", "--yes", action="store_true", help="net apply: do not ask (keep the change if new connections work)")
    ap.add_argument("--programs", default="", help="deploy: programs to install too (docker,images)")
    ap.add_argument("--no-kernel", action="store_true", help="deploy: modules (and programs) only, no kernel upload: no reboot")
    a = ap.parse_args(argv)

    if a.action == "status":
        s, _, slot, trial = hello(a.host, a.port)
        s.close()
        print(f"running slot {slot}" + (" (trial boot, not confirmed)" if trial else ""))
        return

    if a.action == "pair":
        pair(a)
        return

    key = a.key or os.environ.get("RUSTROS_KEY")
    if not key and os.path.exists(a.key_file):
        key = open(a.key_file).read().strip()
    if not key:
        sys.exit("no key: use --key, $RUSTROS_KEY or --key-file")
    key = key.encode()

    if a.action == "deploy":
        sys.exit(deploy(a, key, a.image or "build"))
    if a.action == "upload":
        if not a.image:
            sys.exit("upload needs an image")
        sys.exit(upload(a, key, a.image, a.wait))
    if a.action == "module-install":
        if not a.image or not a.module:
            sys.exit(f"{a.action} needs NAME and a module file")
        module = open(a.module, "rb").read()
        if module[:4] not in (b"\0asm", b"\x7fELF"):
            sys.exit(f"{a.module} is neither a WebAssembly module nor a native module")
        name = a.image.encode()
        reply, _, _ = command(a.host, a.port, key, b"W", bytes([len(name)]) + name + module)
        print(reply)
        sys.exit(0 if reply == "OK" else 1)
    if a.action == "migrate":
        if not a.image:
            sys.exit("migrate needs the PE kernel (build/rustros.efi) and optionally the loader (build/efiboot.efi)")
        sys.exit(migrate(a, key, a.image, a.module or "build/efiboot.efi", a.size_mb))
    if a.action in ("esp-get", "esp-put"):
        if not a.image:
            sys.exit(f"{a.action} needs a FILE")
        sys.exit(esp(a, key))
    if a.action == "run":
        if not a.image:
            sys.exit("run needs a program file (and its arguments)")
        args = [w for w in [a.module, *a.value] if w]
        # an absolute path that is no file here: a program on the machine
        on = a.image.startswith("/") and not os.path.exists(a.image)
        if a.tty:
            size = os.get_terminal_size(sys.stdin.fileno()) if sys.stdin.isatty() else os.terminal_size((80, 24))
            st = interactive(Term(a.host, a.port, key, a.image, args, size.lines, size.columns, a.net, on_machine=on))
            print(f"\r\n[ended: status {st:#x}]" if st is not None else "\r\n[hung up]")
            sys.exit(0 if st == 0 else 1)
        reply, _, _ = command(a.host, a.port, key, b"G", program(a.image, args, a.net, on_machine=on))
        print(reply)
        sys.exit(0 if reply.startswith("OK") else 1)
    if a.action == "inst":
        sys.exit(inst(a, key))
    if a.action == "docker":
        tty = sys.stdin.isatty() and sys.stdout.isatty()
        size = os.get_terminal_size(sys.stdin.fileno()) if tty else os.terminal_size((200, 50))
        t = Term(a.host, a.port, key, "/usr/lib/rustros/docker", dargs, size.lines, size.columns, on_machine=True)
        st = interactive(t) if tty else collect(t)
        if st is None:
            sys.exit("\r\n[hung up]")
        # its exit code, as docker's
        sys.exit((st >> 8) & 0xFF if st & 0x7F == 0 else 128 + (st & 0x7F))
    if a.action == "program-install":
        if not a.image or not a.module:
            sys.exit("program-install needs NAME and a program file")
        data = open(a.module, "rb").read()
        if data[:4] != b"\x7fELF":
            sys.exit(f"{a.module} is not an ELF program")
        name = a.image.encode()
        reply, _, _ = command(a.host, a.port, key, b"Y", bytes([len(name)]) + name + data)
        print(reply)
        sys.exit(0 if reply == "OK" else 1)
    if a.action in ("ps", "kill"):
        words = [a.action] + [w for w in [a.image] if w]
        reply, _, _ = command(a.host, a.port, key, b"J", " ".join(words).encode())
        print(reply, end="" if reply.endswith("\n") else "\n")
        sys.exit(1 if reply.startswith("ERR") else 0)
    if a.action == "log":
        words = [w for w in [a.image, a.module, *a.value] if w]
        reply, _, _ = command(a.host, a.port, key, b"L", " ".join(words).encode())
        print(reply, end="" if reply.endswith("\n") else "\n")
        sys.exit(1 if reply.startswith("ERR") else 0)
    if a.action == "mod":
        words = [w for w in [a.image, a.module, *a.value] if w]
        reply, _, _ = command(a.host, a.port, key, b"N", " ".join(["mod", *words]).encode())
        print(reply, end="" if reply.endswith("\n") else "\n")
        sys.exit(1 if reply.startswith("ERR") else 0)
    if a.action == "net":
        words = [w for w in [a.image, a.module, *a.value] if w]
        if not words:
            words = ["ip"]
        if words[0] == "apply":
            sys.exit(net_apply(a, key, int(words[1]) if len(words) > 1 else 10))
        reply, _, _ = command(a.host, a.port, key, b"N", " ".join(words).encode())
        print(reply)
        sys.exit(1 if reply.startswith("ERR") else 0)
    if a.action == "module-config":
        if not a.image:
            sys.exit("module-config needs NAME [KEY [VALUE | -]]")
        if a.module is None:
            reply, _, _ = command(a.host, a.port, key, b"Q", a.image.encode(), encrypt=True)
            print(reply)
            return
        payload = f"{a.image}\n{a.module}"
        value = " ".join(a.value)
        if value and value != "-":
            payload += "\n" + value
        elif not value:
            sys.exit("module-config NAME KEY needs a VALUE, or '-' to remove it")
        reply, _, _ = command(a.host, a.port, key, b"K", payload.encode(), encrypt=True)
        print(reply)
        sys.exit(0 if reply == "OK" else 1)
    if a.action in ("module-remove", "module-switch"):
        if not a.image:
            sys.exit(f"{a.action} needs NAME")
        kind = a.module or ""
        if a.action == "module-switch" and kind not in ("native", "wasm"):
            sys.exit("module-switch needs native or wasm")
        cmd = b"S" if a.action == "module-switch" else b"X"
        reply, _, _ = command(a.host, a.port, key, cmd, f"{a.image} {kind}".strip().encode())
        print(reply)
        sys.exit(0 if reply == "OK" else 1)

    reply, _, _ = command(a.host, a.port, key, a.action[0].upper().encode())
    print(reply)
    ok = not reply.startswith("ERR") if a.action == "info" else reply == "OK"
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()