#!/usr/bin/env python3
"""Walk `alloy install` to the end, in a VM, without a person at the keyboard.

The installer is a ratatui program with no headless mode and no answer file, so
a scripted install has to answer the six questions the way a person would. What
it must not do is answer them by sending scancodes at the qemu monitor and
screendumping to find out what happened: reading the screen as TEXT is what
makes each step's completion a fact rather than a delay, so every wait below is
for a title the installer only draws once it is on that step, and a slow
install and a stuck one look different. Pictures cannot answer that.

## Two ways in

`--via ssh` (the default) uses the medium's own headless flow:
`alloy-installer-ssh.service` creates the `installer` account on the live medium
only, and the sshd drop-in makes that account's session BE `alloy install` under
a real tty. It is the route a person installing a screenless machine takes, so
it is the one worth exercising.

It could not install anything until 2026-08-26. `alloy install` does not
escalate, and that session landed unprivileged:

    error: Installing to disk: Querying root privilege: This command must be
    executed as the root user
    wipefs: error: /dev/vda: probing initialization failed: Permission denied

Found by running this script (GO alloy `2cf04f20`), and fixed by running the
wizard under `run0` with a polkit grant for the one action run0 asks for. The
Containerfile asserts the three files still agree; if that assertion ever fails,
this is the route that stops working.

`--via serial` boots GRUB's debug entry instead, which carries `alloy.debug` and
so starts `alloy-debug-shell@ttyS0`: a root bash on the guest's serial console.
It was the default while the ssh route was broken, and it is worth keeping for
the case that comes back, and for a medium built without a baked pubkey.

The answers are the ones a regression test wants and not the ones a person
would pick:

  encryption OFF   The unlock path is a whole second thing to go wrong, and
                   nothing downstream of here touches LUKS. `--encrypt` turns
                   it back on for a run that means to exercise it.
  a pubkey         The installed machine has to be reachable afterwards or
                   nothing can be asserted about it.

Use:

    install_drive.py --key ~/.ssh/id_ed25519
    install_drive.py --via serial --pubkey "$(cat ~/.ssh/id_ed25519.pub)"

Exits 0 when the install reports success, 1 when it reports failure, and 3 when
the run could not get far enough to have a verdict. The screen is printed on
every path, because a failed install's own error is the thing worth reading and
it is on it.
"""

import argparse
import os
import re
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from tui import Session, SocketSession, Timeout  # noqa: E402
from vm import Monitor  # noqa: E402

SCRATCH = os.environ.get(
    "VM_STATE", os.path.join(os.path.dirname(os.path.abspath(__file__)), "state")
)

ENTER = "\r"
TAB = "\t"
SPACE = " "


def ssh_argv(key, user, port):
    return [
        "ssh", "-tt",
        # The medium is rebuilt constantly and always answers on the same
        # forwarded port, so a known_hosts entry here is a guaranteed false
        # alarm rather than a check.
        "-o", "StrictHostKeyChecking=no",
        "-o", "UserKnownHostsFile=/dev/null",
        "-o", "LogLevel=ERROR",
        "-o", "ConnectTimeout=10",
        "-p", str(port), "-i", key, "%s@127.0.0.1" % user,
    ]


def connect_ssh(key, port, timeout, first_step=r"step \d+ of 6"):
    """Keep trying until the medium's sshd is up, or give up saying so.

    A live ISO takes tens of seconds to reach multi-user, and the account this
    logs in as is created by a unit inside that boot, so early attempts are
    expected to fail. What is not expected is all of them failing, which is
    what the deadline turns into a message.

    `first_step` is what counts as the wizard having drawn itself. It is a
    pattern rather than `step 1 of 6` because a medium carrying an answer sheet
    opens on the first question it cannot answer, which is not step 1; a caller
    testing that asserts the number itself once the screen is up."""
    deadline = time.time() + timeout
    last = ""
    while time.time() < deadline:
        s = Session(ssh_argv(key, "installer", port),
                    log_path=os.path.join(SCRATCH, "install-drive.raw"))
        try:
            s.wait_for(first_step, min(45.0, max(5.0, deadline - time.time())))
            return s
        except Timeout as e:
            last = e.screen
            s.close()
            time.sleep(3)
    raise SystemExit(
        "error: the installer never drew its first step on port %d.\n"
        "The medium has to carry a baked pubkey matching the key given here -- "
        "build it with `--build-arg ALLOY_SSH_KEY=\"$(cat <key>.pub)\"` -- and "
        "alloy-installer-ssh.service only creates the account when the kernel "
        "command line carries `alloy.installer`, which only the ISO's GRUB "
        "entries set.\nThe last thing the session showed:\n\n%s\n" % (port, last)
    )


# One GRUB menu row: the box border, then the mark, then the title. Keyed on
# the titles this ISO writes rather than on a row number, and the mark is
# matched AFTER the border on purpose. `line.strip()` is not enough: it leaves
# the `\u2502` in front, so `startswith("*")` is false on the selected row and a
# navigator built on it walks to the bottom of the menu and stays there.
GRUB_ROW = re.compile(r"^[^\w*]*(\*?)\s*((?:Install Alloy|Initramfs shell).*?)\s*$")


def grub_menu(s):
    """The menu as `(selected, title)` pairs, or an empty list if it is not up."""
    rows = []
    for line in s.screen.text().splitlines():
        hit = GRUB_ROW.match(line)
        if hit:
            rows.append((hit.group(1) == "*", hit.group(2)))
    return rows


def grub_select(s, entry, timeout=120.0):
    """Stop GRUB's countdown, move to the entry whose title contains `entry`,
    and boot it.

    GRUB draws its menu on the serial console (`make-iso.sh` puts
    `console=ttyS0` on every entry) and marks the current row. So which entry
    is selected is readable, and this navigates by reading rather than by
    counting keystrokes from an assumed starting row.

    That is the fix for two of the things build/vmtest/README.md records as
    costing an afternoon: the countdown expiring before the harness is ready,
    and pressing return without knowing what is highlighted. Any keypress stops
    the countdown, so the loop below holds the menu open by pressing one, and
    nothing is committed until the mark is where it should be.
    """
    monitor = Monitor()
    end = time.time() + timeout
    # Down-then-up is a pair that leaves the selection where it found it, which
    # is what makes holding the menu open safe to do before it has been read.
    while not grub_menu(s):
        if time.time() >= end:
            raise SystemExit("error: GRUB's menu never appeared on the serial console.\n\n%s\n"
                             % s.screen.text())
        monitor.sendkey("down")
        monitor.sendkey("up")
        s.pump(0.3)

    titles = [title for _, title in grub_menu(s)]
    want = next((i for i, title in enumerate(titles) if entry in title), None)
    if want is None:
        raise SystemExit("error: no GRUB entry matching %r. The menu reads:\n  %s\n"
                         % (entry, "\n  ".join(titles)))

    for _ in range(len(titles) * 3 + 6):
        marks = [i for i, (sel, _) in enumerate(grub_menu(s)) if sel]
        if marks == [want]:
            break
        if not marks:
            # The mark is repainted a moment after the move; wait for it rather
            # than pressing again, which is how a navigator overshoots.
            s.pump(0.4)
            continue
        monitor.sendkey("down" if marks[0] < want else "up")
        s.pump(0.4)
    else:
        raise SystemExit("error: could not put GRUB's mark on %r.\n\n%s\n"
                         % (entry, s.screen.text()))

    monitor.sendkey("ret")
    s.pump(1.0)


# Split so the echoed command line does not contain the string being waited
# for. Without that, the wait is satisfied by the terminal echoing the command
# back before the shell has run it, and the next thing typed lands in a shell
# that is not ready.
READY = "VMTEST-SHELL-UP"
READY_CMD = 'echo VMTEST-SHELL"-"UP\n'


def connect_serial(rows, cols, timeout):
    """Boot the debug entry, take its root shell, and start the installer in it.

    The debug entry rather than the default one, and root rather than the ssh
    account, because the ssh account cannot install: see the module docstring.
    `alloy-debug-shell@ttyS0` is gated on `alloy.debug`, which only these GRUB
    entries set, so this reaches nothing on an installed machine."""
    s = SocketSession(os.path.join(SCRATCH, "serial.sock"), rows=rows, cols=cols,
                      log_path=os.path.join(SCRATCH, "install-drive.raw"))
    grub_select(s, "root shell on tty9")

    # The shell arrives some tens of seconds into the boot, and prompts differ,
    # so this asks rather than matching a prompt.
    end = time.time() + timeout
    while True:
        s.send("\n" + READY_CMD, settle=1.0)
        if READY in s.screen.text():
            break
        if time.time() >= end:
            raise SystemExit("error: no root shell on the serial console.\n\n%s\n"
                             % s.screen.text())

    # The size is the guest's, not an ioctl from this end: a serial console has
    # no window size to inherit and defaults to 80x24, which is narrower than
    # the installer's own 80-column budget leaves room for.
    s.send("stty rows %d cols %d; export TERM=xterm-256color\n" % (rows, cols), settle=1.0)
    s.send("clear; alloy install\n", settle=2.0)
    s.wait_for(r"step 1 of 6", 60)
    return s


MARKER = MARKER = "\u25b6"  # alloy_tui::selection::MARKER, the selected-row gutter mark.


def selected_row(s):
    """The name on the row the cursor is on, read off the gutter marker.

    AlloyList draws `MARKER` on the selected row and a space on every other, so
    which row is selected is on the screen rather than something to be counted.
    That is what makes the walk below self-correcting: it checks where it
    landed instead of assuming a starting index and an ordering."""
    hit = re.search(re.escape(MARKER) + r"\s+(\S+)", s.screen.text())
    return hit.group(1) if hit else None


def pick_disk(s, want, limit=40):
    """Put the cursor on `want`.

    `k` to the top first, because the cursor clamps rather than wrapping
    (alloy_tui::Cursor::move_by) so a long enough run lands on row 0 from
    anywhere, and because AlloyList derives its scroll offset from the
    selection — at the top, what is on screen starts at the first disk.

    Then down one row at a time, reading the marker after each. Stepping and
    checking rather than jumping to a counted index is what keeps this honest
    about a list longer than the pane."""
    s.send("k" * limit, settle=0.8)
    seen = []
    for _ in range(limit):
        here = selected_row(s)
        if here is None:
            raise SystemExit(
                "error: the disk step is showing no selected row.\n\n%s\n" % s.screen.text()
            )
        if here == want:
            return here
        if seen and here == seen[-1]:
            break  # the bottom: the cursor clamped and stopped moving
        seen.append(here)
        s.send("j", settle=0.4)
    raise SystemExit(
        "error: no disk called %r on the disk step. The rows it stepped through: %s\n\n%s\n"
        % (want, ", ".join(seen) or "(none)", s.screen.text())
    )


def drive(s, args):
    # 1 of 6, the disk.
    pick_disk(s, args.disk)
    s.send(ENTER)

    # 2 of 6, the machine's name. Enter submits from the name field; the
    # timezone checkbox below it has a default and is left at it.
    s.wait_for(r"step 2 of 6", 30)
    s.send(args.hostname)
    s.send(ENTER)

    # 3 of 6, the account. Tab moves between the four fields, and the pubkey is
    # the last of them.
    s.wait_for(r"step 3 of 6", 30)
    s.send(args.user)
    s.send(TAB)
    s.send(args.password)
    s.send(TAB)
    s.send(args.password)
    s.send(TAB)
    s.send(args.pubkey)
    s.send(ENTER)

    # 4 of 6, encryption. The checkbox has focus and defaults to on, so a space
    # here is what turns it off.
    s.wait_for(r"step 4 of 6", 30)
    if not args.encrypt:
        s.send(SPACE)
    else:
        s.send(TAB)
        s.send(args.passphrase)
        s.send(TAB)
        s.send(args.passphrase)
    s.send(ENTER)

    # 5 of 6, the review. It no longer installs; it advances.
    s.wait_for(r"step 5 of 6", 30)
    s.send(ENTER)

    # 6 of 6, the credits, which is the step that acts. The first Enter raises
    # the erase confirmation and the second answers it — the footer says
    # `enter install` and means it two keystrokes from now.
    s.wait_for(r"step 6 of 6", 30)
    s.send(ENTER)
    # The modal's own sentence, not its footer. `enter confirm` is what any
    # modal's footer reads, so waiting on that would also be satisfied by a
    # different one — and on this screen the next Enter erases a disk.
    s.wait_for(r"Erase .* and install Alloy", 30)
    s.send(ENTER)


def wait_for_verdict(s, seconds):
    """Wait for the run to finish, and say which way.

    Read off the footer, which is the installer's own statement about its
    state: the run screen offers `esc done` only once the sequence is over, and
    adds `r reboot` only when it succeeded (InstallView::hints). So the two
    outcomes are distinguishable without matching on any message text."""
    s.wait_for(r"esc\s+done", seconds)
    # One more frame, so a `reboot` hint drawn in the same repaint as `done`
    # is not missed by reading between the two writes.
    s.pump(1.0)
    return "reboot" in s.screen.text()


def main():
    p = argparse.ArgumentParser(description="drive `alloy install` to the end")
    p.add_argument("--via", choices=("ssh", "serial"), default="ssh",
                   help="ssh: the medium's own headless installer account (the default; "
                        "needs a medium whose baked pubkey matches --key). serial: GRUB's "
                        "debug entry and its root shell")
    p.add_argument("--rows", type=int, default=40)
    p.add_argument("--cols", type=int, default=120)
    p.add_argument("--key", default=os.path.expanduser("~/.ssh/id_ed25519"),
                   help="--via ssh: private key matching the pubkey baked into the medium")
    p.add_argument("--pubkey", default=None,
                   help="public key for the account being created (default: --key + .pub)")
    p.add_argument("--port", type=int, default=2222)
    p.add_argument("--disk", default="vda")
    p.add_argument("--hostname", default="alloytest")
    p.add_argument("--user", default="tester")
    p.add_argument("--password", default="alloytest")
    p.add_argument("--encrypt", action="store_true",
                   help="leave encryption on, and answer its passphrase")
    p.add_argument("--passphrase", default="alloytestalloytest")
    p.add_argument("--connect-timeout", type=float, default=300.0)
    p.add_argument("--install-timeout", type=float, default=1800.0)
    args = p.parse_args()

    if args.pubkey is None:
        path = args.key + ".pub"
        if not os.path.exists(path):
            raise SystemExit("error: no %s; pass --pubkey" % path)
        args.pubkey = open(path).read().strip()

    if args.via == "ssh":
        s = connect_ssh(args.key, args.port, args.connect_timeout, r"step 1 of 6")
    else:
        s = connect_serial(args.rows, args.cols, args.connect_timeout)
    try:
        drive(s, args)
        print("==> installing; this is the long part", flush=True)
        ok = wait_for_verdict(s, args.install_timeout)
        print(s.screen.text())
        if not ok:
            print("\nerror: the installer finished without offering a reboot, "
                  "which is how it says the install failed", file=sys.stderr)
            return 1
        return 0
    except Timeout as e:
        print("error: %s" % e, file=sys.stderr)
        return 3
    finally:
        # Left running, the ForceCommand session would keep the installer alive
        # on a medium the scenario is about to power off.
        s.close()


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