#!/usr/bin/python3 -s
# SPDX-License-Identifier: GPL-2.0-or-later
#
# txnupd-rootfs-setup - prepare a NEW btrfs root file system for transactional
# updates with tukit
#
# Converts /etc of the target root into a nested btrfs subvolume (so it is
# snapshotted together with the rest of the root), and adjusts the target's
# /etc/fstab accordingly. This is meant to be run exactly once by an installer
# or an image build, against a target that is NOT the running system.

import difflib
import os
import re
import shutil
import subprocess
import sys

PROG = "txnupd-rootfs-setup"
ETC_BIND_LINE = "/etc /etc none bind,rw,x-initrd.mount 0 0"
PRESET_UNIT = "create-dirs-from-rpmdb.service"

USAGE = f"""\
Usage: {PROG} [--root DIR] [--dry-run] [--assume-not-live]
       {PROG} --print-fstab FILE

Prepare the btrfs root file system mounted at DIR (default: /) for
transactional updates:

  * convert DIR/etc into a nested btrfs subvolume (copy, verify, then swap)
  * add the "/etc /etc none bind,rw,x-initrd.mount 0 0" entry to DIR/etc/fstab
  * change the "ro" option of the root entry to "ro=vfs"
  * mark the /var and /root entries x-initrd.mount
  * enable create-dirs-from-rpmdb.service (through a preset file in
    DIR/etc/systemd/system-preset, so that the first boot keeps it enabled)
  * if DIR/boot/grub2/grubenv is on btrfs and has no env_block entry, add
    it with grub2-editenv, so that GRUB can save its state (boot_success,
    next_entry, ...) on btrfs

The running system's root is never modified. Use this on a freshly installed
target from an installer, or inside the image build chroot.

Options:
  --root DIR           root of the target file system (default: /)
  --dry-run            check everything and print what would be done
  --assume-not-live    skip the check that DIR is not the running root (needed
                       when /proc/1/root is not readable, e.g. unprivileged)
  --print-fstab FILE   print FILE as it would be changed and exit; touches
                       nothing
  -h, --help           show this help
"""


class Error(Exception):
    """A fatal error; reported as "error: ..." and exits with status 1."""


def info(msg):
    print(f"{PROG}: {msg}", flush=True)


def warn(msg):
    print(f"{PROG}: warning: {msg}", file=sys.stderr, flush=True)


# --- fstab -------------------------------------------------------------------

class Entry:
    """One line of an fstab. Lines that are not entries are kept verbatim."""

    def __init__(self, text):
        self.text = text
        self.fields = None
        stripped = text.lstrip(" \t")
        if not stripped.startswith("#"):
            fields = text.split()
            if len(fields) >= 4:
                self.fields = fields

    @property
    def is_entry(self):
        return self.fields is not None

    @property
    def device(self):
        return self.fields[0]

    @property
    def target(self):
        return self.fields[1]

    @property
    def fstype(self):
        return self.fields[2]

    @property
    def options(self):
        return self.fields[3].split(",")

    @options.setter
    def options(self, opts):
        self.fields[3] = ",".join(opts)

    def render(self):
        # A line that was not changed keeps its original spacing.
        return self.text


def parse_fstab(text):
    lines = text.split("\n")
    if lines and lines[-1] == "":
        lines.pop()
    return [Entry(line) for line in lines]


def has_target(entries, target):
    return any(e.is_entry and e.target == target for e in entries)


def etc_is_bind(entries):
    return any(e.is_entry and e.target == "/etc" and "bind" in e.options
               for e in entries)


def transform_fstab(entries):
    """Rewrite an fstab; returns the new lines (without line endings).

      - "ro" in the options of "/" becomes "ro=vfs"
      - "x-systemd.growfs" is dropped from "/" if another writable mount of the
        same btrfs has it (growfs fails through the read-only mount)
      - an /etc bind mount gets an explicit "rw"
      - "x-initrd.mount" is added to the options of "/var" and "/root"

    Nothing else is touched, including comments and unrelated lines.
    """
    # Devices that have a writable btrfs mount, other than /, that grows the
    # file system.
    growers = {
        e.device for e in entries
        if e.is_entry and e.target != "/" and e.fstype == "btrfs"
        and "x-systemd.growfs" in e.options and "ro" not in e.options
    }

    out = []
    for e in entries:
        if not e.is_entry:
            out.append(e.text)
            continue
        before = e.options
        opts = list(before)

        if e.target == "/" and "ro" in opts:
            opts = ["ro=vfs" if o == "ro" else o for o in opts]

        # systemd-growfs can not resize a btrfs through a read-only mount (it
        # fails with EROFS on every boot). Another, writable, mount of the same
        # file system with x-systemd.growfs does it, so drop the option here,
        # but only when such a mount exists.
        if (e.target == "/" and e.fstype == "btrfs"
                and "x-systemd.growfs" in opts and "ro=vfs" in opts
                and e.device in growers):
            opts = [o for o in opts if o != "x-systemd.growfs"]

        # A bind mount inherits the read-only flag of the mount it is made
        # from (the root is ro=vfs). Without an explicit "rw" it stays
        # read-only, which breaks everything that writes to /etc before a
        # remount: systemd's first boot preset-all, for one.
        if (e.target == "/etc" and "bind" in opts
                and "rw" not in opts and "ro" not in opts):
            new = []
            for o in opts:
                new.append(o)
                if o == "bind":
                    new.append("rw")
            opts = new

        if e.target in ("/var", "/root") and "x-initrd.mount" not in opts:
            opts.append("x-initrd.mount")

        if opts != before:
            e.options = opts
            out.append(" ".join(e.fields))
        else:
            out.append(e.text)
    return out


def wanted_fstab(text):
    """The fstab as it should be, as a list of lines."""
    entries = parse_fstab(text)
    lines = transform_fstab(entries)
    if not has_target(entries, "/etc"):
        lines.append(ETC_BIND_LINE)
    return lines


# --- helpers -----------------------------------------------------------------

def read_text(path):
    with open(path, encoding="utf-8", errors="surrogateescape") as f:
        return f.read()


def fs_type(path):
    r = subprocess.run(["stat", "-f", "-c", "%T", path],
                       capture_output=True, text=True)
    return r.stdout.strip() if r.returncode == 0 else ""


def is_btrfs_subvolume(path):
    # The root directory of every btrfs subvolume has inode number 256.
    return fs_type(path) == "btrfs" and os.stat(path).st_ino == 256


def run(cmd, **kw):
    """Run a command; a failure is fatal."""
    try:
        return subprocess.run(cmd, check=True, **kw)
    except subprocess.CalledProcessError as e:
        raise Error(f"{cmd[0]} failed with status {e.returncode}") from None


def check_not_live(directory, assume_not_live):
    """Fail if DIR is the root of the running system."""
    if assume_not_live:
        return
    if not os.access("/proc/1/root/.", os.R_OK):
        raise Error(
            f"cannot tell whether {directory} is the running system's root "
            "(/proc/1/root is not readable); run as root or pass "
            "--assume-not-live if you are sure")
    a, b = os.stat(directory), os.stat("/proc/1/root/.")
    if (a.st_dev, a.st_ino) == (b.st_dev, b.st_ino):
        raise Error(
            f"{directory} is the root of the running system. This tool only "
            "prepares a new, not yet booted target (run it from the "
            "installer or the image build chroot)")


# --- the work ----------------------------------------------------------------

RSYNC_OPTS = ["--archive", "--hard-links", "--acls", "--xattrs",
              "--numeric-ids", "--one-file-system"]
# rsync itemizes directories whose only difference is a time stamp like this
TIME_ONLY = re.compile(r"\.d\.\.t\.\.\.\.\.\.")


class Setup:
    def __init__(self, root, dry_run, assume_not_live, grub_editenv):
        self.root = root
        self.dry_run = dry_run
        self.assume_not_live = assume_not_live
        self.grub_editenv = grub_editenv
        self.cleanup_new = None

    def path(self, *parts):
        return os.path.join(self.root, *parts)

    # -- /etc --

    def convert_etc(self):
        etc = self.path("etc")
        new = self.path("etc.txnupd-new")
        old = self.path("etc.txnupd-old")

        info("creating subvolume for /etc...")
        self.cleanup_new = new
        run(["btrfs", "subvolume", "create", new], stdout=subprocess.DEVNULL)

        info("copying /etc...")
        run(["rsync", *RSYNC_OPTS, etc + "/", new + "/"])

        # Make sure nothing was missed before the original is touched. A dry
        # run that would still transfer anything means the copy is incomplete.
        r = run(["rsync", *RSYNC_OPTS, "--dry-run", "--itemize-changes",
                 "--delete", etc + "/", new + "/"],
                stdout=subprocess.PIPE, text=True, errors="surrogateescape")
        missed = [line for line in r.stdout.split("\n")
                  if line and not TIME_ONLY.match(line)]
        if missed:
            print("\n".join(missed), file=sys.stderr)
            raise Error("copy of /etc differs from the original, "
                        "nothing was changed")

        info("switching to the new /etc...")
        os.rename(etc, old)
        os.rename(new, etc)
        self.cleanup_new = None
        try:
            shutil.rmtree(old)
        except OSError:
            info(f"warning: could not remove the old copy {old}, "
                 "remove it manually")

    def cleanup(self):
        if self.cleanup_new and os.path.exists(self.cleanup_new):
            subprocess.run(["btrfs", "subvolume", "delete", self.cleanup_new],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # -- fstab --

    def update_fstab(self):
        fstab = self.path("etc", "fstab")
        tmp = self.path("etc", "fstab.txnupd-new")
        if not os.path.isfile(fstab):
            raise Error(f"{fstab} does not exist")

        lines = wanted_fstab(read_text(fstab))
        # cp -a keeps owner, mode and the SELinux label; the content is then
        # written into that file, so the new file has the same attributes.
        run(["cp", "-a", fstab, tmp])
        with open(tmp, "w", encoding="utf-8", errors="surrogateescape") as f:
            f.write("".join(line + "\n" for line in lines))
        os.replace(tmp, fstab)

    # -- create-dirs-from-rpmdb.service --

    def enable_create_dirs(self):
        if not os.path.exists(self.path("usr/lib/systemd/system", PRESET_UNIT)):
            info(f"{PRESET_UNIT} is not installed in the target, "
                 "not enabling it (install tukit)")
            return
        # systemd runs preset-all on the first boot, which would undo a plain
        # "enable" because Fedora's default preset is "disable *". Preset
        # files are read in file name order and the first match wins, so this
        # one has to sort before 90-default.preset.
        presetdir = self.path("etc/systemd/system-preset")
        os.makedirs(presetdir, exist_ok=True)
        with open(os.path.join(presetdir, "00-txnupd-rootfs.preset"), "w") as f:
            f.write(f"enable {PRESET_UNIT}\n")
        run(["systemctl", f"--root={self.root}", "preset", PRESET_UNIT],
            stdout=subprocess.DEVNULL)

    # -- grubenv --
    #
    # GRUB can not write a file that lives on btrfs, and logs "sparse file not
    # allowed" on every boot when it tries to save its environment. Fedora's
    # grub2 can keep the environment block in the reserved area of the btrfs
    # device instead, which is switched on by an "env_block=" entry in grubenv.
    # An image build creates grubenv outside of btrfs, so the entry is missing;
    # any write with grub2-editenv to the file on btrfs adds it (and
    # initialises the block).
    #
    # grub2-editenv only looks for the external block when it is given exactly
    # /boot/grub2/grubenv (not "//boot/grub2/grubenv", and not a path below
    # some other root), so this only works with the target being the root the
    # tool runs in, as in an image build's chroot.

    def grubenv_file(self):
        return self.root.rstrip("/") + "/boot/grub2/grubenv"

    @staticmethod
    def _has_env_block(data):
        return re.search(rb"^env_block=", data, re.M) is not None

    def grubenv_needs_block(self):
        """True if grubenv is a regular file on btrfs without an env_block."""
        f = self.grubenv_file()
        if not os.path.isfile(f) or os.path.islink(f):
            return False
        if fs_type(f) != "btrfs":
            return False
        with open(f, "rb") as fh:
            return not self._has_env_block(fh.read())

    def setup_grubenv(self):
        f = self.grubenv_file()
        if not self.grubenv_needs_block():
            return
        if shutil.which(self.grub_editenv) is None:
            warn(f"{self.grub_editenv} not found, GRUB will not be able to "
                 "save its environment on btrfs")
            return
        # Write back a variable the file already has: nothing changes but the
        # env_block entry.
        with open(f, "rb") as fh:
            m = re.search(rb"^[A-Za-z_][A-Za-z0-9_]*=.*$", fh.read(), re.M)
        if m is None:
            warn(f"{f} has no variables, not adding an env_block entry")
            return
        var = os.fsdecode(m.group(0))
        info(f"adding env_block to {f}...")
        r = subprocess.run([self.grub_editenv, f, "set", var])
        with open(f, "rb") as fh:
            added = self._has_env_block(fh.read())
        if r.returncode != 0 or not added:
            warn(f"could not add an env_block entry to {f}, GRUB will not be "
                 "able to save its environment on btrfs (run the tool in the "
                 "target's chroot with /proc and /dev mounted, or run "
                 "`grub2-editenv /boot/grub2/grubenv set <variable>=<value>` "
                 "in the installed system)")

    # -- driver --

    def run(self):
        root = self.root
        etc = self.path("etc")
        fstab = self.path("etc", "fstab")

        # Preconditions: everything is checked before anything is changed.
        for tool in ("btrfs", "rsync", "stat", "systemctl", "findmnt"):
            if shutil.which(tool) is None:
                raise Error(f"required tool not found: {tool}")

        check_not_live(root, self.assume_not_live)

        if fs_type(etc) != "btrfs":
            raise Error(f"{etc} is not on a btrfs file system")
        if not os.path.isfile(fstab):
            raise Error(f"{fstab} does not exist")

        if (os.path.lexists(self.path("etc.txnupd-new"))
                or os.path.lexists(self.path("etc.txnupd-old"))):
            raise Error(
                f"found {self.path('etc.txnupd-new')} or "
                f"{self.path('etc.txnupd-old')} left over from an "
                "interrupted run; check them and remove manually")

        already_subvol = is_btrfs_subvolume(etc)
        if subprocess.run(["findmnt", "-n", "-M", etc],
                          capture_output=True).returncode == 0:
            raise Error(f"{etc} is a mount point; expected a plain directory")

        text = read_text(fstab)
        entries = parse_fstab(text)
        if has_target(entries, "/etc") and not etc_is_bind(entries):
            raise Error(f"{fstab} has an /etc entry that is not a bind mount; "
                        "refusing to touch it")

        wanted = wanted_fstab(text)
        if (already_subvol and wanted == [e.text for e in entries]
                and not self.grubenv_needs_block()):
            info("/etc is already a subvolume, fstab and grubenv are up to "
                 "date, nothing to do")
            return

        if self.dry_run:
            self.show_dry_run(already_subvol, text, wanted)
            return

        if already_subvol:
            info("/etc is already a subvolume")
        else:
            self.convert_etc()
        self.update_fstab()
        self.enable_create_dirs()
        self.setup_grubenv()
        info(f"done. Relabel {self.path('etc')} (e.g. with setfiles) if the "
             "target uses SELinux and files were created after the last "
             "relabel.")

    def show_dry_run(self, already_subvol, text, wanted):
        info(f"dry run, target: {self.root}")
        if already_subvol:
            info("/etc is already a subvolume, would only update fstab")
        else:
            info(f"would convert {self.path('etc')} into a btrfs subvolume "
                 "(copy, verify, swap)")
        info("fstab would change as follows:")
        fstab = self.path("etc", "fstab")
        old = [line + "\n" for line in text.split("\n")]
        if old and old[-1] == "\n":
            old.pop()
        new = [line + "\n" for line in wanted]
        sys.stdout.writelines(difflib.unified_diff(old, new, fstab, fstab))
        info("would enable create-dirs-from-rpmdb.service with a preset file "
             "(if installed)")
        if self.grubenv_needs_block():
            info(f"would add an env_block entry to {self.grubenv_file()}")


# --- command line ------------------------------------------------------------

def main(argv):
    root = "/"
    dry_run = False
    assume_not_live = False
    print_fstab = None

    args = list(argv)
    while args:
        a = args.pop(0)
        if a == "--root":
            if not args:
                raise Error("--root needs an argument")
            root = args.pop(0)
        elif a.startswith("--root="):
            root = a[len("--root="):]
        elif a == "--dry-run":
            dry_run = True
        elif a == "--assume-not-live":
            assume_not_live = True
        elif a == "--print-fstab":
            if not args:
                raise Error("--print-fstab needs an argument")
            print_fstab = args.pop(0)
        elif a in ("-h", "--help"):
            sys.stdout.write(USAGE)
            return 0
        else:
            sys.stderr.write(USAGE)
            raise Error(f"unknown argument: {a}")

    if print_fstab is not None:
        if not os.path.isfile(print_fstab):
            raise Error(f"{print_fstab}: no such file")
        for line in wanted_fstab(read_text(print_fstab)):
            print(line)
        return 0

    if root.endswith("/"):
        root = root[:-1]
    if not root:
        root = "/"
    if not os.path.isdir(os.path.join(root, "etc")):
        raise Error(f"{root}/etc is not a directory")
    root = os.path.realpath(root)

    setup = Setup(root, dry_run, assume_not_live,
                  os.environ.get("GRUB_EDITENV") or "grub2-editenv")
    try:
        setup.run()
    finally:
        setup.cleanup()
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv[1:]))
    except Error as e:
        print(f"{PROG}: error: {e}", file=sys.stderr)
        sys.exit(1)
    except BrokenPipeError:
        sys.exit(1)
