#!/usr/bin/python3 -I
"""Prepare a dedicated Debian 12 kiosk for the next boot, with local rollback.

Requires the installed application and an existing desktop user's account.
The account, PIN, settings, groups and running desktop are left intact.
"""

import argparse
import base64
from contextlib import contextmanager
import fcntl
import json
import os
from pathlib import Path
import pwd
import re
import stat
import subprocess
import sys
import tempfile


MAX_FILE = 128 * 1024
BACKUP = "/var/lib/4web-kiosk/system-setup"
RECEIPT = BACKUP + "/receipt.json"
ALIAS = "/etc/systemd/system/display-manager.service"
DM_FILE = "/etc/X11/default-display-manager"
CONFIG = "/etc/greetd/config.toml"
DROPIN = "/etc/systemd/system/greetd.service.d/4web-kiosk.conf"
RUNTIME_MASK = "/run/systemd/system/greetd.service"
SNAPSHOTS = (ALIAS, DM_FILE, CONFIG, DROPIN)
UNIT = re.compile(r"[a-zA-Z0-9_][a-zA-Z0-9@_.:-]{0,127}\.service\Z")
USER = re.compile(r"[a-z_][a-z0-9_-]{0,31}\Z")
STATES = {"enabled", "enabled-runtime", "disabled", "static", "indirect", "alias", "masked", "masked-runtime", "not-found"}
RESTART_POLICY = b"[Unit]\nStartLimitIntervalSec=0\n\n[Service]\nRestart=always\nRestartSec=3\n"


class SetupError(Exception):
    pass


class SystemSetup:
    # Alternative roots/owners exist only for isolated unit tests. The CLI
    # always uses the real root filesystem and requires root ownership.
    def __init__(self, root=Path("/"), runner=subprocess.run, lookup=pwd.getpwnam, owner=0, group=0):
        self.root = Path(root).resolve(strict=True)
        self.runner, self.lookup = runner, lookup
        self.owner, self.group = owner, group

    def path(self, name):
        return self.root / name.lstrip("/")

    def sync_parent(self, path):
        descriptor = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
        try:
            os.fsync(descriptor)
        finally:
            os.close(descriptor)

    @contextmanager
    def lock(self):
        path = self.path("/run/4web-kiosk-system.lock")
        self.directory(path.parent)
        descriptor = os.open(path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK, 0o600)
        try:
            info = os.fstat(descriptor)
            if not stat.S_ISREG(info.st_mode) or info.st_uid != self.owner or info.st_mode & 0o077:
                raise SetupError("Unsichere Sperrdatei für die Systemeinrichtung.")
            try:
                fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError:
                raise SetupError("Eine Systemeinrichtung läuft bereits.") from None
            yield
        finally:
            os.close(descriptor)

    def directory(self, path, create=False, mode=0o755):
        if path != self.root:
            self.directory(path.parent, create=create)
        try:
            info = path.lstat()
        except FileNotFoundError:
            if not create:
                return
            path.mkdir(mode=mode)
            os.chown(path, self.owner, self.group, follow_symlinks=False)
            self.sync_parent(path)
            info = path.lstat()
        if (not stat.S_ISDIR(info.st_mode) or info.st_uid != self.owner
                or info.st_mode & 0o022):
            raise SetupError("Unsicheres Systemverzeichnis; Einrichtung abgebrochen.")

    def read(self, path, limit=MAX_FILE):
        self.directory(path.parent)
        descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
        with os.fdopen(descriptor, "rb") as source:
            info = os.fstat(source.fileno())
            if (not stat.S_ISREG(info.st_mode) or info.st_uid != self.owner
                    or info.st_mode & 0o022 or info.st_size > limit):
                raise SetupError("Unsichere oder zu große Systemdatei; Einrichtung abgebrochen.")
            contents = source.read(limit + 1)
            if len(contents) > limit:
                raise SetupError("Systemdatei ist zu groß.")
        return contents, stat.S_IMODE(info.st_mode)

    def snapshot(self, name):
        path = self.path(name)
        self.directory(path.parent)
        try:
            info = path.lstat()
        except FileNotFoundError:
            return {"type": "absent"}
        if stat.S_ISLNK(info.st_mode):
            if name != ALIAS or info.st_uid != self.owner:
                raise SetupError("Unerwarteter Link in einer Systemkonfiguration.")
            target = os.readlink(path)
            if (len(target) > 4096 or any(ord(c) < 32 for c in target)
                    or not UNIT.fullmatch(Path(target).name)):
                raise SetupError("Ungültiger Anmeldemanager-Link.")
            resolved = path.resolve(strict=True)
            if not resolved.is_relative_to(self.root):
                raise SetupError("Anmeldemanager-Link verlässt das Systemverzeichnis.")
            self.read(resolved)
            return {"type": "link", "target": target}
        contents, mode = self.read(path)
        return {"type": "file", "mode": mode, "data": base64.b64encode(contents).decode("ascii")}

    def atomic(self, name, contents, mode=0o644):
        path = self.path(name)
        self.directory(path.parent, create=True)
        if path.exists() or path.is_symlink():
            self.read(path)
        descriptor, temporary = tempfile.mkstemp(prefix=".4web-kiosk-", dir=path.parent)
        try:
            with os.fdopen(descriptor, "wb") as target:
                os.fchown(target.fileno(), self.owner, self.group)
                os.fchmod(target.fileno(), mode)
                target.write(contents)
                target.flush()
                os.fsync(target.fileno())
            os.replace(temporary, path)
            self.sync_parent(path)
        finally:
            Path(temporary).unlink(missing_ok=True)

    def link(self, name, target, allowed_existing=False):
        path = self.path(name)
        self.directory(path.parent, create=True)
        if path.exists() or path.is_symlink():
            if allowed_existing and path.is_symlink() and path.lstat().st_uid == self.owner:
                pass
            else:
                self.read(path)
        descriptor, temporary = tempfile.mkstemp(prefix=".4web-link-", dir=path.parent)
        os.close(descriptor)
        os.unlink(temporary)
        try:
            os.symlink(target, temporary)
            os.lchown(temporary, self.owner, self.group)
            os.replace(temporary, path)
            self.sync_parent(path)
        finally:
            Path(temporary).unlink(missing_ok=True)

    def remove(self, name, allow_link=False):
        path = self.path(name)
        self.directory(path.parent)
        if not path.exists() and not path.is_symlink():
            return
        if not (allow_link and path.is_symlink() and path.lstat().st_uid == self.owner):
            self.read(path)
        path.unlink()
        self.sync_parent(path)

    def command(self, arguments, capture=False, accepted=(0,), **options):
        result = self.runner(arguments, check=False, text=True,
                             stdout=subprocess.PIPE if capture else None,
                             stderr=subprocess.PIPE if capture else None, **options)
        if result.returncode not in accepted:
            raise SetupError("Systembefehl fehlgeschlagen. Die Ausgabe prüfen; anschließend gegebenenfalls --restore ausführen.")
        return (result.stdout or "").strip() if capture else result

    def systemctl(self, *arguments, **options):
        return self.command(["/usr/bin/systemctl", *arguments], **options)

    def check_system(self):
        if sys.platform != "linux" or os.geteuid() != 0:
            raise SetupError("Diesen installierten Befehl als Administrator auf Debian 12 ausführen.")
        release_path = self.path("/etc/os-release")
        if release_path.is_symlink():
            resolved = release_path.resolve(strict=True)
            if not resolved.is_relative_to(self.root):
                raise SetupError("Ungültige Betriebssystem-Konfiguration.")
            release_path = resolved
        release = self.read(release_path)[0].decode("utf-8")
        values = dict(line.split("=", 1) for line in release.splitlines() if "=" in line)
        if values.get("ID", "").strip('"') != "debian" or values.get("VERSION_ID", "").strip('"') != "12":
            raise SetupError("Der Systemkiosk benötigt Debian 12 (bookworm).")
        if not self.path("/run/systemd/system").is_dir():
            raise SetupError("Der Systemkiosk benötigt ein laufendes systemd.")
        for name in ("/usr/bin/4web-kiosk-session", "/usr/lib/4web-kiosk/kiosk-supervisor"):
            if not self.read(self.path(name))[1] & 0o111:
                raise SetupError("Der installierte Sitzungsstarter ist nicht ausführbar.")
        self.directory(self.path("/etc/systemd/system"))

    def user(self, name):
        if not isinstance(name, str) or not USER.fullmatch(name):
            raise SetupError("Einen vorhandenen normalen Linux-Benutzer angeben.")
        try:
            account = self.lookup(name)
        except KeyError:
            raise SetupError("Der angegebene Linux-Benutzer existiert nicht.") from None
        if (account.pw_name != name or type(account.pw_uid) is not int or not 1000 <= account.pw_uid < 65534
                or not isinstance(account.pw_dir, str) or not account.pw_dir.startswith("/")
                or any(ord(c) < 32 for c in account.pw_dir) or len(account.pw_dir) > 4096
                or ".." in Path(account.pw_dir).parts):
            raise SetupError("Das Kiosk-Konto muss ein vorhandener normaler Benutzer mit eigenem Home sein.")
        home = self.path(account.pw_dir)
        self.directory(home.parent)
        info = home.lstat()
        if not stat.S_ISDIR(info.st_mode) or info.st_uid != account.pw_uid or info.st_mode & 0o022:
            raise SetupError("Das Home muss dem Kiosk-Benutzer gehören und darf kein Link oder fremd beschreibbar sein.")
        return account

    def state(self, unit):
        state = self.systemctl("is-enabled", unit, capture=True, accepted=(0, 1, 3, 4))
        # Several systemd releases describe a missing unit only on stderr.
        # No enablement exists to restore for an uninstalled package.
        state = state or "not-found"
        if state not in STATES:
            raise SetupError("Der bisherige Dienstzustand ist unbekannt; keine Änderung vorgenommen.")
        return state

    def baseline(self, account):
        snapshots = {name: self.snapshot(name) for name in SNAPSHOTS}
        previous = self.systemctl("show", "--property=Id", "--value", "display-manager.service", capture=True, accepted=(0, 1, 4))
        if previous == "display-manager.service" or not previous:
            alias = snapshots[ALIAS]
            previous = Path(alias["target"]).name if alias["type"] == "link" else ""
        if previous and not UNIT.fullmatch(previous):
            raise SetupError("Der bisherige Anmeldemanager konnte nicht sicher bestimmt werden.")
        if previous == "greetd.service":
            previous = ""
        target = self.systemctl("get-default", capture=True)
        if not re.fullmatch(r"[a-zA-Z0-9_][a-zA-Z0-9_.@:-]{0,127}\.target", target):
            raise SetupError("Ungültiges bisheriges Boot-Ziel.")
        enabled = {unit: self.state(unit) for unit in {"greetd.service", previous} if unit}
        if enabled["greetd.service"] in {"masked", "masked-runtime"}:
            raise SetupError("greetd ist bereits gesperrt. Die vorhandene Administratorsperre zuerst prüfen.")
        for name in ("/etc/systemd/system/greetd.service", RUNTIME_MASK):
            if self.path(name).exists() or self.path(name).is_symlink():
                raise SetupError("Eine eigene greetd-Unit überschreibt die sichere temporäre Sperre; bitte administrativ prüfen.")
        return dict(format=1, status="prepared", user=account.pw_name, uid=account.pw_uid,
                    home=account.pw_dir, previous=previous, target=target,
                    enabled=enabled, snapshots=snapshots)

    def receipt(self):
        path = self.path(RECEIPT)
        if not path.exists() and not path.is_symlink():
            return None
        directory = path.parent.lstat()
        if directory.st_mode & 0o077:
            raise SetupError("Das Sicherungsverzeichnis muss privat sein (0700).")
        contents, mode = self.read(path, MAX_FILE * 6)
        if mode & 0o077:
            raise SetupError("Die Systemsicherung muss private Dateirechte haben (0600).")
        receipt = json.loads(contents)
        if (not isinstance(receipt, dict) or receipt.get("format") != 1
                or receipt.get("status") not in {"prepared", "enabled", "restored"}
                or not USER.fullmatch(receipt.get("user", ""))
                or not isinstance(receipt.get("snapshots"), dict) or set(receipt["snapshots"]) != set(SNAPSHOTS)
                or not re.fullmatch(r"[a-zA-Z0-9_][a-zA-Z0-9_.@:-]{0,127}\.target", receipt.get("target", ""))
                or not isinstance(receipt.get("enabled"), dict)
                or set(receipt["enabled"]) != ({"greetd.service", receipt.get("previous")} - {""})
                or any(not UNIT.fullmatch(unit) or state not in STATES for unit, state in receipt["enabled"].items())):
            raise SetupError("Die Systemsicherung ist ungültig; nichts wurde zurückgeschrieben.")
        for name, snapshot in receipt["snapshots"].items():
            if not isinstance(snapshot, dict) or snapshot.get("type") not in {"absent", "file", "link"}:
                raise SetupError("Ungültiger Dateieintrag in der Systemsicherung.")
            if snapshot["type"] == "link":
                target = snapshot.get("target")
                if name != ALIAS or not isinstance(target, str) or len(target) > 4096 or not UNIT.fullmatch(Path(target).name) or any(ord(c) < 32 for c in target):
                    raise SetupError("Ungültiger Link in der Systemsicherung.")
            if snapshot["type"] == "file":
                if type(snapshot.get("mode")) is not int or snapshot["mode"] & ~0o755 or snapshot["mode"] & 0o022:
                    raise SetupError("Ungültige Dateirechte in der Systemsicherung.")
                data = base64.b64decode(snapshot.get("data", ""), validate=True)
                if len(data) > MAX_FILE:
                    raise SetupError("Systemsicherung enthält eine zu große Datei.")
        return receipt

    def save(self, receipt):
        self.directory(self.path(BACKUP), create=True, mode=0o700)
        if self.path(BACKUP).stat().st_mode & 0o077:
            raise SetupError("Das Sicherungsverzeichnis muss privat sein (0700).")
        self.atomic(RECEIPT, (json.dumps(receipt, sort_keys=True, indent=2) + "\n").encode(), 0o600)

    def check(self, name):
        self.check_system()
        account = self.user(name)
        receipt = self.receipt()
        if receipt is None:
            self.baseline(account)
        elif receipt["user"] != name or receipt.get("uid") != account.pw_uid:
            raise SetupError("Die vorhandene Sicherung gehört zu einem anderen Kiosk-Konto.")
        return account, receipt

    def install_greetd(self, offline_bundle=None, manifest_sha256=None):
        environment = {**os.environ, "DEBIAN_FRONTEND": "noninteractive", "LC_ALL": "C.UTF-8"}
        if offline_bundle is None:
            self.command(["/usr/bin/apt-get", "install", "-y", "--no-remove", "--no-install-recommends",
                          "greetd", "cage", "dbus-user-session", "procps", "util-linux"], env=environment)
            return
        source = self.read(self.path("/usr/lib/4web-kiosk/offline-install.py"), 512 * 1024)[0]
        namespace = {"__name__": "4web_installed_offline", "__file__": "/usr/lib/4web-kiosk/offline-install.py"}
        exec(compile(source, namespace["__file__"], "exec"), namespace)
        code = namespace.get("ROOT_INSTALL_CODE")
        if not isinstance(code, str):
            raise SetupError("Der installierte Offline-Installer ist unvollständig.")
        self.command(["/usr/bin/python3", "-I", "-c", code, str(offline_bundle), manifest_sha256, "greetd"], env=environment)

    def enable(self, name, offline_bundle=None, manifest_sha256=None):
        self.check_system()
        with self.lock():
            return self.enable_locked(name, offline_bundle, manifest_sha256)

    def enable_locked(self, name, offline_bundle=None, manifest_sha256=None):
        if (offline_bundle is None) != (manifest_sha256 is None):
            raise SetupError("Offline-Ordner und Manifest-Prüfsumme gemeinsam angeben.")
        if offline_bundle is not None:
            offline_bundle = Path(offline_bundle)
            if (not offline_bundle.is_absolute() or offline_bundle.is_symlink() or not offline_bundle.is_dir()
                    or not isinstance(manifest_sha256, str) or not re.fullmatch(r"[0-9a-f]{64}", manifest_sha256)):
                raise SetupError("Ungültiger Offline-Auftrag; absoluter Ordner und SHA-256 sind erforderlich.")
        account, receipt = self.check(name)
        if receipt is not None and receipt["status"] == "enabled":
            return False
        if receipt is not None and receipt["status"] == "restored":
            raise SetupError("Die Einrichtung wurde bereits zurückgenommen. Die alte Systemsicherung vor einer neuen Einrichtung administrativ archivieren.")
        receipt = receipt or self.baseline(account)
        self.save(receipt)  # The immutable baseline is durable before any APT.
        masked = False
        try:
            # Mark the attempted publication first: a directory fsync can fail
            # after the mask has already been atomically renamed into place.
            masked = True
            self.link(RUNTIME_MASK, "/dev/null")
            self.systemctl("daemon-reload")
            self.install_greetd(offline_bundle, manifest_sha256)
            self.read(self.path("/usr/sbin/greetd"), 16 * 1024 * 1024)
            config = ("[terminal]\nvt = 7\n\n[general]\nsource_profile = false\n\n"
                      "[default_session]\ncommand = \"/usr/bin/4web-kiosk-session --locked\"\n"
                      f'user = "{account.pw_name}"\n').encode()
            self.atomic(CONFIG, config)
            self.atomic(DROPIN, RESTART_POLICY)
            self.atomic(DM_FILE, b"/usr/sbin/greetd\n")
            self.remove(RUNTIME_MASK, allow_link=True)
            masked = False
            self.systemctl("daemon-reload")
            if receipt["previous"]:
                self.systemctl("disable", "--no-reload", receipt["previous"])
            self.systemctl("enable", "--no-reload", "--force", "greetd.service")
            vendor = next((name for name in ("/usr/lib/systemd/system/greetd.service", "/lib/systemd/system/greetd.service")
                           if self.path(name).is_file()), None)
            if vendor is None:
                raise SetupError("Die installierte greetd-Servicevorlage fehlt.")
            self.read(self.path(vendor).resolve(strict=True))
            self.link(ALIAS, vendor, allowed_existing=True)
            self.systemctl("set-default", "--no-reload", "graphical.target")
            receipt["status"] = "enabled"
            self.save(receipt)
        except BaseException:
            receipt["status"] = "prepared"
            try:
                if masked:
                    self.remove(RUNTIME_MASK, allow_link=True)
            finally:
                # Even an I/O failure during mask cleanup must not skip
                # restoring the previous display manager and boot target.
                self.restore_saved(receipt, mark_restored=False, reload_runtime=True)
            raise
        return True

    def restore_saved(self, receipt, mark_restored=True, reload_runtime=False):
        # Change only the next boot's links. Reloading while the active display
        # manager's alias points elsewhere can reassign an existing unit's
        # runtime identity. Its process and session must remain untouched.
        self.systemctl("disable", "--no-reload", "greetd.service", accepted=(0, 1, 4))
        for unit, state in receipt["enabled"].items():
            if state in {"enabled", "enabled-runtime", "alias"}:
                arguments = ["enable", "--no-reload", "--force"] + (["--runtime"] if state == "enabled-runtime" else []) + [unit]
                self.systemctl(*arguments)
            elif state in {"masked", "masked-runtime"}:
                self.systemctl("mask", "--no-reload", *(["--runtime"] if state == "masked-runtime" else []), unit)
            else:
                self.systemctl("disable", "--no-reload", unit, accepted=(0, 1, 4))
        for name, snapshot in receipt["snapshots"].items():
            if snapshot["type"] == "absent":
                self.remove(name, allow_link=name == ALIAS)
            elif snapshot["type"] == "link":
                self.link(name, snapshot["target"], allowed_existing=True)
            else:
                if name == ALIAS:
                    self.remove(name, allow_link=True)
                self.atomic(name, base64.b64decode(snapshot["data"], validate=True), snapshot["mode"])
        self.systemctl("set-default", "--no-reload", receipt["target"])
        # A failed enable restores its original alias before refreshing the
        # temporary runtime mask. Normal restore takes effect after reboot.
        if reload_runtime:
            self.systemctl("daemon-reload")
        if mark_restored:
            receipt["status"] = "restored"
        self.save(receipt)

    def restore(self):
        self.check_system()
        with self.lock():
            return self.restore_locked()

    def restore_locked(self):
        receipt = self.receipt()
        if receipt is None:
            raise SetupError("Es gibt keine Systemsicherung für eine Wiederherstellung.")
        if receipt["status"] == "restored":
            return False
        self.restore_saved(receipt)
        return True


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    modes = parser.add_mutually_exclusive_group(required=True)
    modes.add_argument("--check", action="store_true", help="Voraussetzungen prüfen; nichts ändern")
    modes.add_argument("--enable", action="store_true", help="Systemkiosk für den nächsten Neustart einrichten")
    modes.add_argument("--restore", action="store_true", help="Gesicherte Systemkonfiguration wiederherstellen")
    parser.add_argument("--user", help="Vorhandener normaler Linux-Benutzer mit den bisherigen PIN-/Browserdaten")
    parser.add_argument("--offline-bundle", type=Path, help="Absoluter Ordner des geprüften Offline-Bundles")
    parser.add_argument("--manifest-sha256", help="Zuvor geprüfte SHA-256 des Offline-Manifests")
    args = parser.parse_args()
    if args.restore and any((args.user, args.offline_bundle, args.manifest_sha256)):
        parser.error("--restore benötigt keine weiteren Angaben")
    if not args.restore and not args.user:
        parser.error("--check und --enable benötigen --user")
    if args.check and any((args.offline_bundle, args.manifest_sha256)):
        parser.error("--check prüft nur die installierten Systemvoraussetzungen")
    try:
        setup = SystemSetup()
        if args.check:
            setup.check(args.user)
            print("Voraussetzungen geprüft. Es wurde nichts verändert.")
        elif args.restore:
            changed = setup.restore()
            print("Vorherige Systemkonfiguration wiederhergestellt." if changed else "Die vorherige Systemkonfiguration ist bereits wiederhergestellt.")
            print("Die Änderung wird beim nächsten manuellen Neustart wirksam.")
        else:
            changed = setup.enable(args.user, args.offline_bundle, args.manifest_sha256)
            print("Systemkiosk für den nächsten Start eingerichtet." if changed else "Der Systemkiosk ist bereits eingerichtet.")
            print("Die laufende Sitzung bleibt offen. Bei Gelegenheit manuell neu starten.")
            print("Rückweg mit Administrator-Anmeldung: sudo 4web-kiosk-setup-system --restore")
        return 0
    except SetupError as error:
        print("Systemkiosk-Einrichtung fehlgeschlagen: " + str(error), file=sys.stderr)
        return 1
    except (OSError, ValueError, TypeError, UnicodeError, subprocess.SubprocessError):
        print("Systemkiosk-Einrichtung fehlgeschlagen. Die Systemkonfiguration und gegebenenfalls --restore als Administrator prüfen.", file=sys.stderr)
        return 1


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