#!/usr/bin/env python3

import subprocess
import logging
import datetime
import os
import re
import sys
import shlex


# =========================
# Einstellungen
# =========================

ZFS_DATASETS = [
    "rpool/ROOT/pve-1",
    "rpool/pveconf"
]

SNAPSHOT_TAG = "pve-update-via-cron"

# Anzahl der ZFS-Snapshots, die je Dataset behalten werden.
MAX_SNAPSHOTS = 5

# Anzahl älterer Kernel, die zusätzlich zum aktuell
# laufenden Kernel als Fallback behalten werden.
#
# 0 = nur aktiven Kernel behalten
# 1 = aktiven + 1 Fallback
# 2 = aktiven + 2 Fallback-Kernel
KEEP_OLD_KERNELS = 1

LOGFILE = "/var/log/proxmox_update.log"

REBOOT_FLAG = os.environ.get("REBOOT", "NO").upper()

TIMESTAMP = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")


# Unterstützt sowohl historische pve-kernel-* Pakete
# als auch aktuelle proxmox-kernel-* Pakete.
KERNEL_PKG_PATTERN = re.compile(
    r"^(?:pve-kernel|proxmox-kernel)-([0-9][^\s]*?)(?:-signed)?$"
)


# Cron-Jobs laufen mit einem minimalen PATH.
# Proxmox-Tools liegen teilweise in /usr/sbin bzw. /sbin.
os.environ["PATH"] = (
    "/usr/sbin:/usr/bin:/sbin:/bin:"
    + os.environ.get("PATH", "")
)


# =========================
# Logging
# =========================

logging.basicConfig(
    filename=LOGFILE,
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)


def log(msg, level="info"):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    print(f"[{timestamp}] {msg}")

    if level == "info":
        logging.info(msg)

    elif level == "warn":
        logging.warning(msg)

    elif level == "error":
        logging.error(msg)


# =========================
# Kommando Helfer
# =========================

def run(cmd, check=True, quiet=False):
    """
    Führt einen Befehl ohne Shell aus.

    cmd kann eine Argumentliste oder ein String ohne
    Shell-Metazeichen sein.

    quiet=True unterdrückt das automatische Loggen
    von stdout.
    """

    if isinstance(cmd, str):
        cmd = shlex.split(cmd)

    log(
        f"Ausführen: {' '.join(cmd)}"
    )

    try:
        result = subprocess.run(
            cmd,
            check=check,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        if result.stdout.strip() and not quiet:
            log(result.stdout.strip())

        if result.stderr.strip() and not check:
            log(result.stderr.strip(), "warn")

        return result.stdout.strip()

    except FileNotFoundError as e:

        log(
            f"FEHLER: Befehl nicht gefunden: "
            f"{' '.join(cmd)} ({e})",
            "error",
        )

        if check:
            raise

        return ""

    except subprocess.CalledProcessError as e:

        log(
            f"FEHLER: {' '.join(cmd)}",
            "error",
        )

        if e.stderr:
            log(
                e.stderr.strip(),
                "error",
            )

        if check:
            raise

        return ""


# =========================
# Kernel Prüfung
# =========================

def get_running_kernel():
    return run([
        "uname",
        "-r",
    ])


def version_key(version):
    """
    Sortierbarer Schlüssel für Kernel-Versionen.

    Beispiel:
    6.14.8
    6.14.10
    6.14.11
    """

    return [
        int(part) if part.isdigit() else part
        for part in re.split(r"[.\-]", version)
    ]


def _query_kernel_packages():
    """
    Fragt dpkg nach beiden Proxmox-Kernel-Paketgruppen ab.

    Zurückgegeben werden Tupel:

        (status, paketname)

    Beispiel:

        ("ii", "proxmox-kernel-6.14.11-1-pve-signed")
        ("rc", "proxmox-kernel-6.8.12-2-pve-signed")

    Nur "ii" bedeutet tatsächlich installiert.
    """

    output = run(
        [
            "dpkg-query",
            "-W",
            "-f=${db:Status-Abbrev} ${Package}\n",
            "pve-kernel-*",
            "proxmox-kernel-*",
        ],
        check=False,
        quiet=True,
    )

    packages = []

    for line in output.splitlines():

        parts = line.split(None, 1)

        if len(parts) != 2:
            continue

        status, pkg = parts

        packages.append(
            (status, pkg)
        )

    return packages


def get_installed_kernel_versions():
    """
    Liefert tatsächlich installierte Proxmox-Kernel.

    Nur dpkg-Status "ii" wird berücksichtigt.
    """

    packages = _query_kernel_packages()

    kernels = []

    for status, pkg in packages:

        if status != "ii":
            continue

        match = KERNEL_PKG_PATTERN.match(pkg)

        if not match:
            continue

        kernels.append({
            "version": match.group(1),
            "package": pkg,
        })

    return kernels


def get_installed_kernel():
    """
    Ermittelt die höchste installierte Proxmox-Kernel-Version.
    """

    kernels = get_installed_kernel_versions()

    if not kernels:
        return ""

    versions = sorted(
        {
            kernel["version"]
            for kernel in kernels
        },
        key=version_key,
    )

    log(
        "Installierte Proxmox-Kernel: "
        + ", ".join(versions)
    )

    return versions[-1]


def cleanup_old_kernels():
    """
    Entfernt alte installierte Proxmox-Kernel.

    Sicherheitsregeln:

    - Der aktuell laufende Kernel wird niemals gelöscht.
    - Ist der laufende Kernel nicht eindeutig in der
      installierten Kernel-Liste vorhanden, wird nichts gelöscht.
    - KEEP_OLD_KERNELS ältere Kernel bleiben erhalten.
    - Ein neuerer installierter Kernel als der laufende
      wird ebenfalls nicht gelöscht, da er nach einem
      Reboot benötigt werden kann.
    - Nur Status "ii" wird hier berücksichtigt.
    - Entfernen erfolgt über apt-get purge.
    """

    running = get_running_kernel()

    if not running:
        log(
            "⚠️ Aktiver Kernel konnte nicht ermittelt werden - "
            "keine Kernel werden gelöscht",
            "warn",
        )
        return

    kernels = get_installed_kernel_versions()

    if not kernels:
        log(
            "⚠️ Keine installierten Proxmox-Kernel erkannt - "
            "keine Kernel werden gelöscht",
            "warn",
        )
        return

    # Version -> Pakete
    version_packages = {}

    for kernel in kernels:

        version = kernel["version"]
        package = kernel["package"]

        version_packages.setdefault(
            version,
            []
        ).append(package)

    installed_versions = sorted(
        version_packages.keys(),
        key=version_key,
    )

    log(
        f"🖥️ Aktiver Kernel: {running}"
    )

    log(
        "📋 Installierte Proxmox-Kernel:"
    )

    for version in reversed(installed_versions):

        marker = ""

        if version == running:
            marker = " ← AKTIV"

        log(
            f"   {version}{marker}"
        )

    # Der laufende Kernel muss exakt als installierte
    # Kernel-Version vorhanden sein.
    active_matches = [
        version
        for version in installed_versions
        if version == running
    ]

    if len(active_matches) != 1:

        log(
            "⚠️ Aktiver Kernel wurde nicht eindeutig in den "
            "installierten Kernel-Paketen gefunden - "
            "keine Kernel werden gelöscht",
            "warn",
        )

        return

    active_index = installed_versions.index(running)

    # Alles unterhalb des laufenden Kernels ist älter.
    older_versions = installed_versions[:active_index]

    # Die direkt älteren KEEP_OLD_KERNELS Versionen bleiben.
    fallback_versions = older_versions[
        max(0, len(older_versions) - KEEP_OLD_KERNELS):
    ]

    for version in reversed(fallback_versions):

        log(
            f"🛡️ Fallback-Kernel bleibt erhalten: {version}"
        )

    # Sicherheitsregel:
    #
    # Versionen oberhalb des laufenden Kernels werden NICHT
    # gelöscht. Das kann ein neuer Kernel sein, auf den nach
    # dem Reboot gewechselt werden soll.
    newer_versions = [
        version
        for version in installed_versions
        if version_key(version) > version_key(running)
    ]

    for version in newer_versions:

        log(
            f"🛡️ Neuerer installierter Kernel bleibt erhalten: "
            f"{version} - Reboot kann erforderlich sein",
            "warn",
        )

    protected_versions = {
        running,
        *fallback_versions,
        *newer_versions,
    }

    remove_versions = [
        version
        for version in installed_versions
        if version not in protected_versions
    ]

    if not remove_versions:

        log(
            "OK: Keine alten Kernel zur Entfernung vorhanden"
        )

        return

    for version in remove_versions:

        packages_to_remove = version_packages.get(
            version,
            [],
        )

        for package in packages_to_remove:

            log(
                f"🗑️ Entferne alten Kernel: {package}"
            )

            result = subprocess.run(
                [
                    "apt-get",
                    "purge",
                    "-y",
                    package,
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )

            if result.returncode == 0:

                log(
                    f"✅ Kernel entfernt: {package}"
                )

            else:

                log(
                    f"❌ Fehler beim Entfernen von {package} "
                    f"(Exit-Code {result.returncode})",
                    "error",
                )

                if result.stderr.strip():

                    log(
                        result.stderr.strip(),
                        "error",
                    )

                # Absichtlich kein raise:
                # Fehler bei einem einzelnen Kernel sollen
                # den restlichen Update-Lauf nicht abbrechen.


def cleanup_old_kernel_configs():
    """
    Entfernt ausschließlich dpkg-Konfigurationsreste ("rc")
    alter Proxmox-Kernel-Pakete.
    """

    packages = _query_kernel_packages()

    residual_pkgs = [
        pkg
        for status, pkg in packages
        if status == "rc"
    ]

    if not residual_pkgs:

        log(
            "OK: Keine Konfigurationsreste alter Kernel-Pakete gefunden"
        )

        return

    residual_pkgs = sorted(residual_pkgs)

    total = len(residual_pkgs)

    log(
        f"🧼 Entferne Konfigurationsreste von "
        f"{total} alten Kernel-Paketen"
    )

    for index, pkg in enumerate(
        residual_pkgs,
        start=1,
    ):

        log(
            f"({index}/{total}) Lösche Reste von: {pkg}"
        )

        result = subprocess.run(
            [
                "apt-get",
                "purge",
                "-y",
                "-qq",
                pkg,
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        if result.returncode == 0:

            log(
                f"✅ ({index}/{total}) Erledigt: {pkg}"
            )

        else:

            log(
                f"❌ ({index}/{total}) Fehler beim Entfernen "
                f"der Reste: {pkg}",
                "error",
            )

            if result.stderr.strip():

                log(
                    result.stderr.strip(),
                    "error",
                )


def reboot_required():
    """
    Prüft getrennt von der Kernelbereinigung, ob ein Reboot
    erforderlich ist.

    Reboot erforderlich, wenn der laufende Kernel nicht
    der höchsten installierten Kernel-Version entspricht.
    """

    running = get_running_kernel()
    installed = get_installed_kernel()

    log(
        f"Kernel aktiv: {running}"
    )

    log(
        f"Kernel installiert: {installed}"
    )

    if not running or not installed:

        log(
            "⚠️ Kernel-Vergleich nicht möglich - "
            "gehe sicherheitshalber von 'kein Reboot nötig' aus",
            "warn",
        )

        return False

    return running != installed


# =========================
# Boot Prüfung
# =========================

def proxmox_boot_ok():

    try:

        result = subprocess.run(
            [
                "proxmox-boot-tool",
                "status",
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

    except FileNotFoundError as e:

        log(
            f"FEHLER: proxmox-boot-tool nicht gefunden: {e}",
            "error",
        )

        return False

    if result.stdout.strip():

        log(
            result.stdout.strip()
        )

    if result.returncode == 0:

        return True

    if result.stderr.strip():

        log(
            result.stderr.strip(),
            "error",
        )

    else:

        log(
            "FEHLER: proxmox-boot-tool status "
            f"lieferte Exit-Code {result.returncode}",
            "error",
        )

    return False


def refresh_boot():

    log(
        "Aktualisiere Proxmox Bootloader"
    )

    run([
        "proxmox-boot-tool",
        "refresh",
    ])


# =========================
# Update Prüfung
# =========================

def has_updates():

    run([
        "apt-get",
        "update",
    ])

    output = run(
        [
            "apt-get",
            "-s",
            "dist-upgrade",
        ],
        check=False,
    )

    updates = [
        line
        for line in output.splitlines()
        if line.startswith("Inst ")
    ]

    if updates:

        log(
            f"{len(updates)} Updates verfügbar"
        )

        return True

    log(
        "OK: Keine Updates verfügbar"
    )

    return False


# =========================
# ZFS Snapshot Funktionen
# =========================

def zfs_dataset_exists(dataset):

    result = subprocess.run(
        [
            "zfs",
            "list",
            "-H",
            "-o",
            "name",
            dataset,
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )

    return result.returncode == 0


def create_snapshot(dataset):

    snapshot = (
        f"{dataset}@{SNAPSHOT_TAG}-{TIMESTAMP}"
    )

    log(
        f"Snapshot erstellen: {snapshot}"
    )

    run([
        "zfs",
        "snapshot",
        snapshot,
    ])


def cleanup_snapshots(dataset):

    output = run(
        [
            "zfs",
            "list",
            "-H",
            "-t",
            "snapshot",
            "-o",
            "name",
            "-s",
            "creation",
        ],
        check=False,
        quiet=True,
    )

    snapshots = []

    prefix = (
        f"{dataset}@{SNAPSHOT_TAG}-"
    )

    for line in output.splitlines():

        if line.startswith(prefix):

            snapshots.append(line)

    if len(snapshots) <= MAX_SNAPSHOTS:

        return

    for snap in snapshots[:-MAX_SNAPSHOTS]:

        log(
            f"Lösche Snapshot: {snap}"
        )

        run([
            "zfs",
            "destroy",
            "-r",
            snap,
        ])


def create_update_snapshots():

    for dataset in ZFS_DATASETS:

        if not zfs_dataset_exists(dataset):

            log(
                f"⚠️ ZFS-Dataset nicht vorhanden: {dataset} - "
                "Snapshot wird übersprungen, Update läuft weiter",
                "warn",
            )

            continue

        cleanup_snapshots(dataset)

        create_snapshot(dataset)


# =========================
# Upgrade
# =========================

def upgrade_system():

    log(
        "⚙️ Starte dist-upgrade"
    )

    run([
        "apt-get",
        "dist-upgrade",
        "-y",
    ])

    log(
        "🧼 Starte autoremove"
    )

    run([
        "apt-get",
        "autoremove",
        "--purge",
        "-y",
    ])

    # Eigene Kernelbereinigung nach dem normalen
    # Debian/Proxmox-Upgrade.
    cleanup_old_kernels()

    # Danach eventuell verbliebene rc-Reste bereinigen.
    cleanup_old_kernel_configs()


# =========================
# Hauptprogramm
# =========================

def main():

    log(
        "===== Proxmox Update gestartet ====="
    )

    log(
        f"REBOOT-Flag: {REBOOT_FLAG}"
    )

    # Exit-Code 2 für nicht-root beibehalten.
    if os.geteuid() != 0:

        log(
            "FEHLER: Muss als root laufen",
            "error",
        )

        sys.exit(2)

    kernel_before = get_running_kernel()

    log(
        f"Kernel vor Update: {kernel_before}"
    )

    if has_updates():

        # Snapshots VOR dem eigentlichen dist-upgrade.
        create_update_snapshots()

        upgrade_system()

        # Nach Kernel-/Paket-Updates Bootloader aktualisieren.
        refresh_boot()

    else:

        log(
            "ℹ️ Keine Updates installiert"
        )

        # Auch ohne verfügbare Updates regelmäßig alte
        # installierte Kernel bereinigen.
        cleanup_old_kernels()

        cleanup_old_kernel_configs()

    # Reboot-Erkennung bleibt vollständig getrennt
    # von der Kernelbereinigung.
    if reboot_required():

        log(
            "⚠️ Neuer Kernel vorhanden - Reboot erforderlich",
            "warn",
        )

        if REBOOT_FLAG == "YES":

            log(
                "Prüfe proxmox-boot-tool"
            )

            if proxmox_boot_ok():

                log(
                    "OK: Bootloader OK"
                )

                log(
                    "REBOOT=YES erkannt - Neustart in 1 Minute"
                )

                run([
                    "shutdown",
                    "-r",
                    "+1",
                    "Proxmox-Update abgeschlossen",
                ])

            else:

                log(
                    "❌ Bootloader Prüfung fehlgeschlagen - "
                    "kein Neustart",
                    "error",
                )

        else:

            log(
                "REBOOT-Flag ist nicht 'YES' - "
                "kein automatischer Neustart"
            )

        # Monitoring muss weiterhin erkennen,
        # dass ein Reboot erforderlich ist.
        sys.exit(1)

    log(
        "OK: Update erfolgreich abgeschlossen"
    )

    sys.exit(0)


if __name__ == "__main__":
    main()