pc-unlocker/app/puk/disks.py

"""Blockgeräte auflisten, mounten und wieder aushängen.

Alle Operationen setzen root-Rechte voraus (Offline-Recovery läuft aus einem
Live-System). NTFS wird über ntfs-3g gemountet; Fast-Startup/Hibernation von
Windows wird mit ``remove_hiberfile`` behandelt, sonst verweigert ntfs-3g das
Schreiben.
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
import tempfile
from dataclasses import dataclass, field


class DirtyVolume(RuntimeError):
    """NTFS konnte nicht RW gemountet werden (Dirty-Bit / Fast-Startup / Hibernation).

    Das Volume laesst sich nach einem ``fix_ntfs()`` in aller Regel wieder
    schreibbar einhaengen.
    """


# Textbausteine, an denen ntfs-3g einen unsauberen Zustand meldet
_DIRTY_MARKERS = (
    "unclean",
    "was not safely closed",
    "dirty",
    "hibernat",
    "metadata kept in windows cache",
    "falling back to read-only",
    "windows is in",
)


@dataclass
class Partition:
    path: str            # z.B. /dev/sda2
    size: str            # menschenlesbar, z.B. "256G"
    fstype: str          # ext4, ntfs, vfat, ...
    label: str           # Datenträgerbezeichnung, kann leer sein
    mountpoint: str      # aktueller Mountpunkt laut lsblk (meist leer)
    our_mount: str = ""  # von uns gesetzter Mountpunkt
    os_kind: str = ""    # "linux", "windows" oder "" (unbekannt/nicht gemountet)

    @property
    def display(self) -> str:
        bits = [self.path, self.size or "?", self.fstype or "?"]
        if self.label:
            bits.append(f"[{self.label}]")
        if self.os_kind:
            bits.append(f"-> {self.os_kind.upper()}")
        return "  ".join(bits)


def list_partitions() -> list[Partition]:
    """Alle mountbaren Partitionen zurueckgeben (ohne Swap/leere FS)."""
    out = subprocess.run(
        ["lsblk", "-J", "-b", "-o",
         "PATH,SIZE,FSTYPE,LABEL,MOUNTPOINT,TYPE"],
        capture_output=True, text=True, check=True,
    ).stdout
    data = json.loads(out)
    parts: list[Partition] = []

    def walk(nodes):
        for n in nodes:
            fstype = n.get("fstype") or ""
            ntype = n.get("type") or ""
            if ntype in ("part", "lvm", "crypt") and fstype not in ("", "swap"):
                parts.append(Partition(
                    path=n.get("path") or "",
                    size=_human(n.get("size")),
                    fstype=fstype,
                    label=n.get("label") or "",
                    mountpoint=n.get("mountpoint") or "",
                ))
            if n.get("children"):
                walk(n["children"])

    walk(data.get("blockdevices", []))
    return parts


def _human(nbytes) -> str:
    try:
        n = int(nbytes)
    except (TypeError, ValueError):
        return ""
    for unit in ("B", "K", "M", "G", "T", "P"):
        if n < 1024 or unit == "P":
            return f"{n:.0f}{unit}" if unit == "B" else f"{n:.1f}{unit}"
        n /= 1024
    return ""


def mount(part: Partition, force: bool = False) -> str:
    """Partition read-write mounten und OS-Typ erkennen. Gibt Mountpunkt zurueck.

    Bei NTFS wird ein unsauberer Zustand (Dirty-Bit / Fast-Startup /
    Hibernation) als :class:`DirtyVolume` gemeldet, damit die Oberflaeche
    nachfragen kann. Mit ``force=True`` wird nach einem ``fix_ntfs()`` ein
    erzwungener Schreib-Mount versucht.
    """
    if part.our_mount and os.path.ismount(part.our_mount):
        return part.our_mount

    mp = tempfile.mkdtemp(prefix="puk-")
    if part.fstype == "ntfs":
        # remove_hiberfile: raeumt hiberfile.sys weg (Fast-Startup/Hibernation).
        # force: Schreib-Mount trotz gesetztem Dirty-Flag erzwingen.
        opts = "rw,remove_hiberfile"
        if force:
            opts += ",force"
        cmd = ["mount", "-t", "ntfs-3g", "-o", opts, part.path, mp]
    else:
        cmd = ["mount", "-o", "rw", part.path, mp]

    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode != 0:
        os.rmdir(mp)
        err = (res.stderr or res.stdout or "").strip()
        if part.fstype == "ntfs" and _looks_dirty(err):
            raise DirtyVolume(err or "NTFS-Volume ist unsauber (Dirty-Bit).")
        raise RuntimeError(err or f"mount fehlgeschlagen ({part.path})")

    part.our_mount = mp
    part.os_kind = detect_os(mp)
    return mp


def _looks_dirty(message: str) -> bool:
    low = message.lower()
    return any(m in low for m in _DIRTY_MARKERS)


def fix_ntfs(part: Partition) -> str:
    """Dirty-Bit loeschen und NTFS-Journal zuruecksetzen (ntfsfix).

    Darf nur auf einem *nicht* gemounteten Volume laufen. Gibt die Ausgabe
    von ntfsfix zurueck.
    """
    if part.our_mount and os.path.ismount(part.our_mount):
        raise RuntimeError("Volume ist noch gemountet -- erst aushaengen.")
    if shutil.which("ntfsfix") is None:
        raise RuntimeError("ntfsfix nicht gefunden (Paket ntfs-3g).")
    res = subprocess.run(["ntfsfix", "-d", part.path],
                         capture_output=True, text=True)
    out = (res.stdout + res.stderr).strip()
    if res.returncode != 0:
        raise RuntimeError(out or "ntfsfix fehlgeschlagen")
    return out


def detect_os(mountpoint: str) -> str:
    if os.path.exists(os.path.join(mountpoint, "etc", "shadow")):
        return "linux"
    # Windows: SAM-Hive, Gross-/Kleinschreibung variiert je nach FS
    for cfg in ("Windows/System32/config", "WINDOWS/system32/config"):
        if os.path.exists(os.path.join(mountpoint, cfg, "SAM")):
            return "windows"
    return ""


def unmount(part: Partition) -> None:
    if part.our_mount and os.path.ismount(part.our_mount):
        subprocess.run(["umount", part.our_mount],
                       capture_output=True, text=True)
    if part.our_mount and os.path.isdir(part.our_mount):
        try:
            os.rmdir(part.our_mount)
        except OSError:
            pass
    part.our_mount = ""
    part.os_kind = ""


def unmount_all(parts: list[Partition]) -> None:
    for p in parts:
        try:
            unmount(p)
        except Exception:
            pass