#!/usr/bin/python3
"""
Copyright (C) 2026 Shawn W Dunn

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.

txnupd-maint — Maintenance utilities for transactional and automatic-update systems.

Wraps snapshot lifecycle and bootloader management operations, making them
available as standalone commands outside of a full transactional-update run.

Subcommands that manage a snapshot transaction:
    shell              Open an interactive shell inside a new snapshot
    run <cmd>          Run a command inside a new snapshot

Subcommands that operate on an already-open snapshot (or the live system):
    update-grub        Regenerate grub.cfg (GRUB2 systems; no-op on BLS)
    update-initrd      Rebuild the initrd / UKI images
    update-bootloader  Reinstall the bootloader

Snapshot lifecycle (Open, Close, Abort) and non-interactive in-snapshot
commands are dispatched through the tukitd D-Bus interface
(org.opensuse.tukit.Transaction) rather than shelling out to the tukit CLI.

The interactive 'shell' and 'run' subcommands still use tukit(8) as a
subprocess because org.opensuse.tukit.Transaction.Call is asynchronous and
captures output into a string — it cannot attach a live TTY for interactive
use.

sdbootutil(8) and the direct-mode bootloader tools (grub2-mkconfig, dracut,
pbl) are always invoked as subprocesses; they have no D-Bus or library
interface.
"""

import argparse
import logging
import os
import shlex
import shutil
import subprocess
import sys

import dbus
import dbus.mainloop.glib
from gi.repository import GLib

try:
    import selinux as _selinux
    _HAVE_SELINUX = True
except ImportError:
    _HAVE_SELINUX = False

LOG = logging.getLogger("txnupd-maint")

# True is substituted by CMake (configure_file @ONLY).
# True on openSUSE builds; False on distributions that do not ship pbl(8).
_WITH_PBL: bool = True

_TUKIT_BUS      = "org.opensuse.tukit"
_TUKIT_TX_PATH  = "/org/opensuse/tukit/Transaction"
_TUKIT_TX_IFACE = "org.opensuse.tukit.Transaction"


# ---------------------------------------------------------------------------
# BLS detection
# ---------------------------------------------------------------------------

def _bls_tool() -> str | None:
    """
    Return the available BLS bootloader management tool, or None (GRUB2 path).

    Preference order:
      'sdbootutil' — openSUSE transactional systems; snapshot-aware ESP management.
      'bootctl'    — any systemd-boot system (e.g. Fedora) without sdbootutil.
      None         — no BLS bootloader detected; callers use the GRUB2/pbl path.
    """
    if shutil.which("sdbootutil") is not None:
        r = subprocess.run(["/usr/bin/sdbootutil", "is-installed"], capture_output=True)
        if r.returncode == 0:
            return "sdbootutil"
    if shutil.which("bootctl") is not None:
        r = subprocess.run(["/usr/bin/bootctl", "is-installed"], capture_output=True)
        if r.returncode == 0:
            return "bootctl"
    return None


# ---------------------------------------------------------------------------
# tukitd D-Bus helpers
# ---------------------------------------------------------------------------

def _tukit_iface() -> dbus.Interface:
    return dbus.Interface(
        dbus.SystemBus().get_object(_TUKIT_BUS, _TUKIT_TX_PATH),
        _TUKIT_TX_IFACE,
    )


def _tukit_open(base_id: str | None) -> str | None:
    """
    Create a new read-write snapshot via tukitd and return its ID.

    Passes base_id to Open(); "active" is used when no base is specified,
    which creates the snapshot from the currently running system.
    """
    try:
        snapshot_id = str(_tukit_iface().Open(base_id or "active"))
        LOG.info("Opened snapshot %s", snapshot_id)
        return snapshot_id
    except dbus.DBusException as exc:
        msg = exc.get_dbus_message()
        if "not activatable" in msg or "not provided" in msg:
            msg += " — is tukitd installed and running?"
        LOG.error("tukit Open failed: %s", msg)
        return None


def _tukit_close(snapshot_id: str) -> int:
    """Finalize and promote a snapshot to the default root filesystem."""
    try:
        _tukit_iface().Close(snapshot_id)
        LOG.info("Snapshot %s promoted to default — reboot to apply changes.", snapshot_id)
        return 0
    except dbus.DBusException as exc:
        LOG.error("tukit Close %s failed: %s", snapshot_id, exc.get_dbus_message())
        return 1


def _tukit_abort(snapshot_id: str) -> int:
    """Discard a snapshot without promoting it."""
    try:
        _tukit_iface().Abort(snapshot_id)
        LOG.info("Snapshot %s aborted.", snapshot_id)
        return 0
    except dbus.DBusException as exc:
        LOG.error("tukit Abort %s failed: %s", snapshot_id, exc.get_dbus_message())
        return 1


def _tukit_call(snapshot_id: str, argv: list[str]) -> int:
    """
    Execute a non-interactive command in an open snapshot via tukitd.

    org.opensuse.tukit.Transaction.Call is asynchronous: the D-Bus method
    returns immediately and the result arrives later via a CommandExecuted
    or Error signal.  A GLib main loop is run until the signal for our
    snapshot ID arrives.

    The command string is built with shlex.join() so that individual
    arguments survive the wordexp() expansion tukitd applies server-side.

    stdout/stderr from the command are captured by tukitd and returned as a
    string in the signal payload; they are logged line-by-line on arrival.

    Not suitable for interactive commands that need a live TTY or stdin —
    use _run_in_transaction() for those.
    """
    command = shlex.join(argv)
    loop = GLib.MainLoop()
    result: dict[str, int | None] = {"rc": None}

    def _on_executed(snapshot, returncode, output):
        if str(snapshot) != snapshot_id:
            return
        result["rc"] = int(returncode)
        for line in str(output).splitlines():
            LOG.info("%s", line)
        loop.quit()

    def _on_error(snapshot, returncode, output):
        if str(snapshot) != snapshot_id:
            return
        result["rc"] = int(returncode) if returncode else 1
        LOG.error("Error in snapshot %s (exit %d): %s", snapshot, result["rc"], output)
        loop.quit()

    bus = dbus.SystemBus()
    bus.add_signal_receiver(
        _on_executed,
        signal_name="CommandExecuted",
        dbus_interface=_TUKIT_TX_IFACE,
        path=_TUKIT_TX_PATH,
    )
    bus.add_signal_receiver(
        _on_error,
        signal_name="Error",
        dbus_interface=_TUKIT_TX_IFACE,
        path=_TUKIT_TX_PATH,
    )

    try:
        dbus.Interface(
            bus.get_object(_TUKIT_BUS, _TUKIT_TX_PATH), _TUKIT_TX_IFACE
        ).Call(snapshot_id, command)
    except dbus.DBusException as exc:
        LOG.error("tukit Call failed: %s", exc.get_dbus_message())
        return 1

    loop.run()
    return result["rc"] if result["rc"] is not None else 1


# ---------------------------------------------------------------------------
# Helpers for running commands in-snapshot or directly
# ---------------------------------------------------------------------------

def _run(
    *args: str,
    snapshot_id: str | None = None,
    in_snapshot: bool = True,
) -> int:
    """
    Run a command, optionally inside an open snapshot.

    When snapshot_id is set and in_snapshot is True, dispatches to
    _tukit_call() which uses the tukitd D-Bus interface.  Otherwise runs
    directly as a subprocess on the live system.

    sdbootutil operations pass in_snapshot=False because that tool runs on
    the host and takes the snapshot ID as a plain argument.
    """
    if snapshot_id and in_snapshot:
        return _tukit_call(snapshot_id, list(args))
    r = subprocess.run(list(args))
    if r.returncode != 0:
        LOG.error("Command failed (exit %d): %s", r.returncode, " ".join(args))
    return r.returncode


def _path_writable(path: str) -> bool:
    """Return True if the filesystem containing path is mounted read-write."""
    try:
        st = os.statvfs(path)
        return not (st.f_flag & os.ST_RDONLY)
    except OSError:
        return True  # can't determine; let the operation try and fail naturally


def _selinux_enabled() -> bool:
    if _HAVE_SELINUX:
        return bool(_selinux.is_selinux_enabled())
    binary = shutil.which("selinuxenabled") or "/usr/sbin/selinuxenabled"
    return subprocess.run([binary], capture_output=True).returncode == 0


# ---------------------------------------------------------------------------
# shell / run subcommands
# ---------------------------------------------------------------------------

def _run_in_transaction(
    argv: list[str],
    base_id: str | None,
    continue_id: str | None,
    abort_on_failure: bool,
) -> int:
    """
    Core transaction lifecycle used by both 'shell' and 'run'.

    Snapshot Open/Close/Abort go through the tukitd D-Bus interface.
    The command itself is executed via 'tukit call' as a subprocess because
    org.opensuse.tukit.Transaction.Call captures output and cannot attach
    a live TTY — both shell and run need real stdin/stdout/stderr.
    """
    if continue_id:
        snapshot_id = continue_id
        LOG.info("Resuming work in existing snapshot %s", snapshot_id)
    else:
        snapshot_id = _tukit_open(base_id)
        if snapshot_id is None:
            return 1

    env = os.environ.copy()
    env["PS1"] = r"txnupd-maint \w# "

    r = subprocess.run(["tukit", "call", snapshot_id, *argv], env=env)
    shell_rc = r.returncode

    if shell_rc == 0:
        return _tukit_close(snapshot_id)

    if abort_on_failure:
        LOG.warning("Command exited %d — aborting snapshot %s.", shell_rc, snapshot_id)
        _tukit_abort(snapshot_id)
    else:
        LOG.warning("Command exited %d — snapshot %s kept open.", shell_rc, snapshot_id)
        LOG.info(
            "To resume: txnupd-maint %s --continue %s",
            "shell" if argv == ["bash"] else "run -- " + " ".join(argv),
            snapshot_id,
        )

    return shell_rc


def cmd_shell(base_id: str | None, continue_id: str | None, abort_on_failure: bool) -> int:
    """Open an interactive bash shell inside a new (or existing) snapshot."""
    return _run_in_transaction(["bash"], base_id, continue_id, abort_on_failure)


def cmd_run(
    command: list[str],
    base_id: str | None,
    continue_id: str | None,
    abort_on_failure: bool,
) -> int:
    """Run an arbitrary command inside a new (or existing) snapshot."""
    if not command:
        LOG.error("run: no command specified.")
        return 2
    return _run_in_transaction(command, base_id, continue_id, abort_on_failure)


# ---------------------------------------------------------------------------
# Bootloader subcommands
# ---------------------------------------------------------------------------

def cmd_update_grub(snapshot_id: str | None, no_reboot: bool) -> int:
    """
    Regenerate grub.cfg.

    On BLS systems (systemd-boot) this is a no-op: the first GRUB
    configuration is embedded in the EFI binary and menu entries are
    generated dynamically by the blscfg command.

    On GRUB2 systems, runs grub2-mkconfig via _tukit_call() when a snapshot
    is specified, or directly otherwise.  Relabels the output file if
    SELinux is enforcing.
    """
    if _bls_tool() is not None:
        LOG.info("BLS bootloader detected — grub.cfg is generated dynamically, nothing to do.")
        return 0

    auto_snapshot = False
    if snapshot_id is None and not _path_writable("/boot/grub2"):
        LOG.info(
            "/boot/grub2 is on a read-only filesystem — "
            "opening a new snapshot to apply the update."
        )
        snapshot_id = _tukit_open(None)
        if snapshot_id is None:
            return 1
        auto_snapshot = True

    LOG.info("Regenerating grub.cfg")
    rc = _run(
        "/usr/sbin/grub2-mkconfig", "--output=/boot/grub2/grub.cfg",
        snapshot_id=snapshot_id,
    )
    if rc != 0:
        if auto_snapshot:
            _tukit_abort(snapshot_id)
        return rc

    if _selinux_enabled():
        LOG.info("Relabeling /boot/grub2/grub.cfg for SELinux")
        if _HAVE_SELINUX:
            try:
                _selinux.restorecon("/boot/grub2/grub.cfg")
            except OSError as exc:
                LOG.warning("restorecon failed: %s", exc)
                rc = 1
        else:
            rc = _run("restorecon", "/boot/grub2/grub.cfg", snapshot_id=None)

    if auto_snapshot:
        if rc == 0:
            return _tukit_close(snapshot_id)
        _tukit_abort(snapshot_id)
        return rc

    if not no_reboot:
        LOG.info("grub.cfg updated — a reboot is recommended.")

    return rc


def cmd_update_initrd(snapshot_id: str | None) -> int:
    """
    Rebuild the initrd / UKI images.

    On openSUSE BLS systems (sdbootutil): delegates to sdbootutil mkinitrd,
    which runs dracut inside the snapshot chroot and installs the result into
    the ESP with snapshot-keyed content-addressable paths.  A snapshot ID is
    required on this path.

    On other BLS systems (bootctl, e.g. Fedora): runs dracut directly; the
    kernel-install machinery handles ESP placement.  No snapshot ID is needed
    or accepted.

    On GRUB2 systems: runs dracut via _tukit_call() when a snapshot is
    specified, or directly otherwise.
    """
    tool = _bls_tool()
    if tool == "sdbootutil":
        if not snapshot_id:
            LOG.error(
                "BLS systems with sdbootutil require --snapshot <ID> for update-initrd "
                "(sdbootutil needs the snapshot ID to install the UKI)."
            )
            return 2
        LOG.info("Rebuilding initrd via sdbootutil (snapshot %s)", snapshot_id)
        return _run("/usr/bin/sdbootutil", "mkinitrd", snapshot_id, snapshot_id=None)
    if tool == "bootctl":
        LOG.info("Rebuilding initrd via dracut")
        return _run("dracut", "--force", "--regenerate-all", snapshot_id=None)

    auto_snapshot = False
    if snapshot_id is None and not _path_writable("/boot"):
        LOG.info(
            "/boot is on a read-only filesystem — "
            "opening a new snapshot to apply the update."
        )
        snapshot_id = _tukit_open(None)
        if snapshot_id is None:
            return 1
        auto_snapshot = True

    LOG.info("Rebuilding initrd via dracut")
    rc = _run("dracut", "--force", "--regenerate-all", snapshot_id=snapshot_id)
    if auto_snapshot:
        if rc == 0:
            return _tukit_close(snapshot_id)
        _tukit_abort(snapshot_id)
    return rc


def cmd_update_bootloader(snapshot_id: str | None) -> int:
    """
    Reinstall the bootloader.

    On openSUSE BLS systems (sdbootutil): delegates to sdbootutil update,
    which copies the bootloader binary from the snapshot's rootfs into the
    ESP if a newer version is available.  pbl is intentionally bypassed for
    BLS because it has only partial BLS support and would incorrectly replace
    the shim (bsc#1228864).  A snapshot ID is required on this path.

    On other BLS systems (bootctl, e.g. Fedora): runs bootctl update, which
    copies the installed systemd-boot binary into the ESP if newer.  No
    snapshot ID is needed or accepted.

    On GRUB2 systems: runs pbl via _tukit_call() when a snapshot is specified,
    or directly otherwise.
    """
    tool = _bls_tool()
    if tool == "sdbootutil":
        if not snapshot_id:
            LOG.error(
                "BLS systems with sdbootutil require --snapshot <ID> for update-bootloader "
                "(sdbootutil needs the snapshot ID)."
            )
            return 2
        LOG.info("Reinstalling bootloader via sdbootutil (snapshot %s)", snapshot_id)
        return _run("/usr/bin/sdbootutil", "update", snapshot_id, snapshot_id=None)
    if tool == "bootctl":
        LOG.info("Updating bootloader via bootctl")
        return _run("/usr/bin/bootctl", "update", snapshot_id=None)

    if _WITH_PBL:
        auto_snapshot = False
        if snapshot_id is None and not _path_writable("/boot/grub2"):
            LOG.info(
                "/boot/grub2 is on a read-only filesystem — "
                "opening a new snapshot to apply the update."
            )
            snapshot_id = _tukit_open(None)
            if snapshot_id is None:
                return 1
            auto_snapshot = True

        LOG.info("Reinstalling bootloader via pbl")
        rc = _run("/sbin/pbl", "--install", snapshot_id=snapshot_id)
        if auto_snapshot:
            if rc == 0:
                return _tukit_close(snapshot_id)
            _tukit_abort(snapshot_id)
        return rc

    LOG.error(
        "update-bootloader on GRUB2 systems requires pbl, "
        "which is not available in this build."
    )
    return 1


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

def _add_transaction_args(p: argparse.ArgumentParser) -> None:
    """Add the shared base/continue/abort flags to a transaction subcommand."""
    group = p.add_mutually_exclusive_group()
    group.add_argument(
        "--base", metavar="ID",
        help="Create the new snapshot based on snapshot ID rather than the current default.",
    )
    group.add_argument(
        "--continue", dest="continue_id", metavar="ID",
        help="Resume work in an existing open snapshot instead of creating a new one.",
    )
    p.add_argument(
        "--abort-on-failure",
        action="store_true",
        help="Discard the snapshot if the command exits non-zero (default: keep it open).",
    )


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="txnupd-maint",
        description="Maintenance utilities for transactional and automatic-update systems.",
    )
    p.add_argument("-v", "--verbose", action="store_true", help="Enable debug logging.")

    sub = p.add_subparsers(dest="command", metavar="COMMAND")
    sub.required = True

    # shell
    ss = sub.add_parser(
        "shell",
        help="Open an interactive shell inside a new snapshot.",
        description=(
            "Creates a new read-write snapshot and opens an interactive bash "
            "shell inside it.  On clean exit the snapshot is promoted to the "
            "new default root.  On non-zero exit it is kept open for resumption "
            "with --continue."
        ),
    )
    _add_transaction_args(ss)

    # run
    sr = sub.add_parser(
        "run",
        help="Run a command inside a new snapshot.",
        description=(
            "Creates a new read-write snapshot and runs the given command inside "
            "it.  On success the snapshot is promoted to the new default root.  "
            "On failure it is kept open (or aborted with --abort-on-failure)."
        ),
    )
    _add_transaction_args(sr)
    sr.add_argument("cmd", nargs=argparse.REMAINDER, metavar="CMD",
                    help="Command and arguments to run inside the snapshot.")

    # update-grub
    sg = sub.add_parser(
        "update-grub",
        help="Regenerate grub.cfg (GRUB2 systems only; no-op on BLS/systemd-boot).",
    )
    sg.add_argument("--snapshot", metavar="ID",
                    help="Run inside this open snapshot chroot rather than the live system.")
    sg.add_argument("--no-reboot", action="store_true",
                    help="Suppress the reboot recommendation message.")

    # update-initrd
    si = sub.add_parser(
        "update-initrd",
        help="Rebuild the initrd (or UKI images on BLS systems).",
    )
    si.add_argument("--snapshot", metavar="ID",
                    help="Snapshot ID to operate on (required for BLS systems).")

    # update-bootloader
    sb = sub.add_parser(
        "update-bootloader",
        help="Reinstall the bootloader.",
    )
    sb.add_argument("--snapshot", metavar="ID",
                    help="Snapshot ID to operate on (required for BLS systems).")

    return p


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(name)s: %(levelname)s %(message)s",
    )

    p = build_parser()
    args = p.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    if os.geteuid() != 0:
        LOG.error("txnupd-maint must be run as root.")
        sys.exit(1)

    # Must be called before any dbus.SystemBus() connection is opened so
    # that D-Bus signal reception integrates with the GLib main loop used
    # by _tukit_call().
    dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)

    match args.command:
        case "shell":
            rc = cmd_shell(args.base, args.continue_id, args.abort_on_failure)
        case "run":
            # argparse REMAINDER includes a leading '--' when the user writes
            # 'txnupd-maint run -- cmd arg'; strip it if present.
            cmd = args.cmd
            if cmd and cmd[0] == "--":
                cmd = cmd[1:]
            rc = cmd_run(cmd, args.base, args.continue_id, args.abort_on_failure)
        case "update-grub":
            rc = cmd_update_grub(args.snapshot, args.no_reboot)
        case "update-initrd":
            rc = cmd_update_initrd(args.snapshot)
        case "update-bootloader":
            rc = cmd_update_bootloader(args.snapshot)
        case _:
            p.print_help()
            rc = 1

    sys.exit(rc)


if __name__ == "__main__":
    main()
