diskos/diskos_installer/state.py

"""Per-user state so the tool can always put the device back to bone-stock and
so `remove` can clean up after itself.

Everything the tool creates lives under ONE directory (state_dir()); nothing is
written system-wide. `remove` deletes that directory (and, if the user asks, the
tool itself) - that's the whole uninstall footprint on Linux/macOS.

Contents:
    state.json            what we've done (installed?, versions, timestamps)
    stock/stock.bin       the saved stock-rootfs image, rebuilt from the user's firmware (for restore-stock)
    stock/stock.sha256    its checksum
    build/                scratch for extract/build (cleaned opportunistically)
"""

import hashlib
import json
import os
import shutil
import sys
import time

from . import errors

APP = "diskos-installer"
# A marker file dropped in our state dir. wipe_all() refuses to recursively delete any directory
# that does not contain it, so a mis-set DISKOS_INSTALLER_HOME can never delete an arbitrary tree.
SENTINEL = ".diskos-installer-state"
# The diskOS install/restore lifecycle values written to state.json["phase"]. Used to recognise a
# legacy (pre-sentinel) state dir by CONTENT, specifically enough that a foreign project's state.json
# cannot masquerade as ours. Keep in sync with service.py.
_KNOWN_PHASES = {"prepared", "flash-started", "flash-verified", "restore-started", "restore-verified"}


def _fsync_dir(path):
    """Fsync a directory so a rename/create inside it is durable across a crash."""
    try:
        fd = os.open(path, os.O_RDONLY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)
    except OSError:
        pass


def _state_base():
    """Compute the state-dir PATH only - no directory creation, no side effects.
    Used by wipe_all() so validation never re-creates the sentinel it is about to check."""
    env = os.environ.get("DISKOS_INSTALLER_HOME")
    if env:
        return env
    if sys.platform == "darwin":
        return os.path.expanduser(f"~/Library/Application Support/{APP}")
    xdg = os.environ.get("XDG_DATA_HOME")
    return os.path.join(xdg, APP) if xdg else os.path.expanduser(f"~/.local/share/{APP}")


def _is_our_statedir(base):
    """STRONG test that `base` is (or was) a diskOS state dir - governs whether we may adopt it and,
    in wipe_all(), recursively delete it. Requires GENUINE evidence, never a merely-common basename,
    so an unrelated directory that happens to contain a 'build/' or a 'state.json' is NEVER matched:
      - our sentinel file, OR
      - our exact stock-recovery pair stock/stock.bin + stock/stock.sha256, OR
      - a state.json that parses AND matches our specific schema (a boolean `installed` plus a
        `phase` drawn from our known lifecycle values) - not merely the generic key NAMES, so an
        unrelated project's state.json cannot be mistaken for ours and later deleted."""
    if os.path.isfile(os.path.join(base, SENTINEL)):
        return True
    if (os.path.isfile(os.path.join(base, "stock", "stock.bin"))
            and os.path.isfile(os.path.join(base, "stock", "stock.sha256"))):
        return True
    sj = os.path.join(base, "state.json")
    if os.path.isfile(sj):
        try:
            with open(sj) as f:
                d = json.load(f)
            if (isinstance(d, dict) and isinstance(d.get("installed"), bool)
                    and isinstance(d.get("phase"), str) and d.get("phase") in _KNOWN_PHASES):
                return True
        except (OSError, ValueError, TypeError):
            pass
    return False


def state_dir():
    """A single, per-user, non-system directory for all tool state (created on first use)."""
    base = _state_base()
    marker = os.path.join(base, SENTINEL)
    # SAFETY, only for an EXPLICITLY-set home: refuse to adopt a pre-existing, non-empty directory
    # that shows NO sign of being ours, so a mis-set DISKOS_INSTALLER_HOME (e.g. $HOME) can't be
    # tagged here and later wiped by 'remove'. The DEFAULT XDG/APP path is always ours; and an
    # existing state dir (even one predating the sentinel) is recognised by its contents.
    if (os.environ.get("DISKOS_INSTALLER_HOME") and os.path.isdir(base)
            and not _is_our_statedir(base)):
        try:
            nonempty = bool(os.listdir(base))
        except OSError:
            nonempty = True
        if nonempty:
            raise errors.PreflightError(
                f"DISKOS_INSTALLER_HOME={base!r} is a non-empty directory that is not a diskOS "
                "state dir - refusing to use it. Point it at a new or empty path.", code="E142")
    os.makedirs(base, exist_ok=True)
    if not os.path.exists(marker):
        try:
            with open(marker, "w") as f:
                f.write("diskOS installer state directory - safe to delete via 'diskos-installer remove'\n")
        except OSError:
            pass
    return base


def _state_path():
    return os.path.join(state_dir(), "state.json")


def load():
    """Return the saved state. A corrupt file is reported (not silently erased) so
    'never installed' and 'state damaged' stay distinguishable."""
    p = _state_path()
    if not os.path.exists(p):
        return {}
    try:
        with open(p) as f:
            return json.load(f)
    except ValueError:
        return {"_corrupt": True}
    except OSError:
        return {}


def save(d):
    d = {k: v for k, v in d.items() if k != "_corrupt"}
    d["updated"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    tmp = _state_path() + ".tmp"
    with open(tmp, "w") as f:
        json.dump(d, f, indent=2)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, _state_path())
    _fsync_dir(state_dir())


def sha256_file(path, progress=None):
    h = hashlib.sha256()
    total = os.path.getsize(path)
    read = 0
    with open(path, "rb") as f:
        while True:
            chunk = f.read(1 << 20)
            if not chunk:
                break
            h.update(chunk)
            read += len(chunk)
            if progress:
                progress(read, total)
    return h.hexdigest()


def stock_paths():
    d = os.path.join(state_dir(), "stock")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, "stock.bin"), os.path.join(d, "stock.sha256")


def save_stock_image(src_bin, progress=None):
    """Copy the saved stock-rootfs flashable image into state so restore-stock can
    always reflash exactly it. Atomic + validated: the previous good image is only
    replaced once the new one is fully written, hashed, magic-checked, and fsynced -
    so a crash/disk-full can't destroy your recovery image. Returns its sha256."""
    dst, shafile = stock_paths()
    d = os.path.dirname(dst)
    tmp_bin, tmp_sha = dst + ".tmp", shafile + ".tmp"

    h = hashlib.sha256()
    total = os.path.getsize(src_bin)
    read = 0
    with open(src_bin, "rb") as s, open(tmp_bin, "wb") as o:
        while True:
            chunk = s.read(1 << 20)
            if not chunk:
                break
            o.write(chunk)
            h.update(chunk)
            read += len(chunk)
            if progress:
                progress(read, total)
        o.flush()
        os.fsync(o.fileno())
    digest = h.hexdigest()

    # validate BEFORE committing (never publish a truncated/non-squashfs image)
    with open(tmp_bin, "rb") as f:
        if f.read(4) != b"hsqs":
            os.unlink(tmp_bin)
            raise ValueError("refusing to save a non-squashfs stock image (bad magic).")
    with open(tmp_sha, "w") as f:
        f.write(digest + "\n")
        f.flush()
        os.fsync(f.fileno())

    # commit: rename sha first, then the image, then fsync the dir. If interrupted
    # between renames, have_stock_image() sees a mismatch and reports 'no valid image'
    # rather than trusting a half-committed pair.
    os.replace(tmp_sha, shafile)
    os.replace(tmp_bin, dst)
    _fsync_dir(d)
    return digest


def stock_image_exists():
    """Fast existence check (no hashing) - safe to call on the UI thread."""
    dst, shafile = stock_paths()
    return os.path.exists(dst) and os.path.exists(shafile)


def have_stock_image():
    dst, shafile = stock_paths()
    try:
        if not (os.path.exists(dst) and os.path.exists(shafile)):
            return False
        want = open(shafile).read().strip()
        return sha256_file(dst) == want
    except OSError:
        return False


def build_dir():
    d = os.path.join(state_dir(), "build")
    os.makedirs(d, exist_ok=True)
    return d


def wipe_all():
    """Delete the entire tool state directory (the full uninstall footprint).
    Returns (path, errors): errors is a list of (path, message) for anything that
    could NOT be removed, so the caller reports the truth instead of a false success."""
    d = _state_base()   # PATH only - never re-create the sentinel we are about to verify
    errs = []
    # SAFETY: only ever recursively delete a directory we can PROVE is our own state dir. This
    # guards against DISKOS_INSTALLER_HOME being set to '/', '$HOME', the cwd, or any other tree.
    real = os.path.realpath(d)
    forbidden = {os.path.realpath(os.sep), os.path.realpath(os.path.expanduser("~"))}
    try:
        forbidden.add(os.path.realpath(os.getcwd()))
    except OSError:
        pass
    if real in forbidden:
        return d, [(d, "refusing to delete a root/home/current directory")]
    if not (os.path.isdir(d) and _is_our_statedir(d)):
        return d, [(d, "not a diskOS state dir (no sentinel or known state) - refusing to delete")]
    try:
        if os.stat(d).st_uid != os.getuid():
            return d, [(d, "state dir is not owned by you - refusing to delete")]
    except OSError as e:
        return d, [(d, str(e))]
    shutil.rmtree(d, onerror=lambda fn, path, exc: errs.append((path, str(exc[1]))))
    return d, errs