pigway-raspberry-pi-agent/app/storage_cache.py

"""Daily filesystem capacity cache, persisted in RAM across monitor restarts."""
import json
import time
from pathlib import Path
from shutil import disk_usage

class StorageInventory:
    """Read host block devices without mounting, probing SMART, or writing disks."""

    def __init__(self, clock=time.time, runner=None, mounts_path="/proc/1/mountinfo", usage=None, cache_path=None, capacity_path=None):
        import subprocess
        self.clock = clock
        self.runner = runner or (lambda: subprocess.check_output(
            ["lsblk", "--json", "--bytes", "--paths", "--output",
             "NAME,KNAME,TYPE,SIZE,FSTYPE,UUID,LABEL,MODEL,SERIAL,MAJ:MIN"],
            text=True, timeout=5, stderr=subprocess.DEVNULL))
        self.mounts_path = Path(mounts_path)
        self.usage = usage or (lambda mount: disk_usage("/proc/1/root" + mount))
        self.next_update_at = 0
        self.observed_at = None
        self.error = False
        self.devices = []
        self.cache_path = Path(cache_path) if cache_path else None
        self.scanned_at = None
        self.capacity_path = Path(capacity_path) if capacity_path else None
        if self.capacity_path:
            try:
                saved=json.loads(self.capacity_path.read_text())
                at=saved['observed_at']
                if isinstance(at,(int,float)) and 0 <= self.clock()-at < 604800 and isinstance(saved['devices'],list):
                    self.devices=saved['devices'];self.observed_at=at
                    self.scanned_at=saved['scanned_at'];self.next_update_at=at+604800
            except (OSError,ValueError,KeyError,TypeError):pass

    def scan(self):
        if self.cache_path is None:
            raise ValueError("storage inventory cache path required")
        try:
            previous = json.loads(self.cache_path.read_text())
            if 0 <= self.clock() - previous["scanned_at"] < 60:
                return {"accepted": True, "scanned_at": previous["scanned_at"], "cooldown": True}
        except (OSError, ValueError, KeyError, TypeError):
            pass
        data = json.loads(self.runner())
        self.discover(data)
        document = {"scanned_at": self.clock(), "inventory": data}
        temporary = self.cache_path.with_suffix(".tmp")
        with open(temporary,"w",opener=lambda path, flags: __import__("os").open(path, flags, 0o600)) as file:
            json.dump(document, file)
        temporary.replace(self.cache_path)
        return {"accepted": True, "scanned_at": document["scanned_at"], "cooldown": False}

    def sample(self, config=None):
        data = None
        changed = False
        if self.cache_path is not None:
            try:
                document = json.loads(self.cache_path.read_text())
                changed = document["scanned_at"] != self.scanned_at
                data = document["inventory"]
                self.scanned_at = document["scanned_at"]
            except FileNotFoundError:
                pass
            except (OSError, ValueError, KeyError, TypeError):
                self.error = True
        if data is not None and (changed or self.clock() >= self.next_update_at):
            try:
                self.devices = self.discover(data)
                self.observed_at = self.clock()
                self.error = False
            except (OSError, ValueError, TypeError, KeyError, __import__("subprocess").SubprocessError):
                self.error = True
            self.next_update_at = self.clock() + (3600 if self.error else 604800)
            if self.capacity_path and not self.error:
                try:
                    self.capacity_path.write_text(json.dumps(dict(devices=self.devices,observed_at=self.observed_at,scanned_at=self.scanned_at)))
                except OSError:pass
        names = dict(config.items("storage_names", raw=True)) if config is not None and config.has_section("storage_names") else {}
        visible = dict(config.items("storage_visible", raw=True)) if config is not None and config.has_section("storage_visible") else {}
        result = []
        for device in self.devices:
            item = dict(device)
            item["name"] = names.get(item["id"], "auto")
            if item["name"].lower() == "auto":
                item["name"] = item["label"] or item["mountpoint"] or item["device"]
            item["visible"] = visible.get(item["id"], "true").lower() != "false"
            if self.error:
                item["state"] = "unknown"
            result.append(item)
        known = {d["id"] for d in result}
        for identifier in sorted((names.keys() | visible.keys()) - known):
            result.append(dict(id=identifier, name=names.get(identifier, identifier),
                               visible=visible.get(identifier, "true").lower() != "false",
                               state="unknown" if self.error else "missing", device="", mountpoint="",
                               label="", capacity_bytes=None, used_bytes=None, free_bytes=None,
                               used_percent=None))
        return result

    def discover(self, data):
        import hashlib
        import re
        mounts, sources = {}, {}
        for line in self.mounts_path.read_text().splitlines():
            left, right = line.split(" - ", 1)
            fields = left.split()
            # Bind mounts and Btrfs subvolumes must not duplicate filesystem capacity.
            mount = re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), fields[4])
            row = (fields[3] != "/", len(mount), mount)
            mounts.setdefault(fields[2], []).append(row)
            source = right.split()[1]
            if source.startswith("/dev/"):
                sources.setdefault(str(Path(source).resolve()), []).append(row)
        devices = data["blockdevices"]
        if not isinstance(devices, list):
            raise ValueError("invalid block device inventory")
        result, seen = [], set()

        def walk(device, serial="", parent=""):
            kind = device.get("type")
            if kind in {"loop", "rom", "zram"}:
                return
            serial = str(device.get("serial") or serial).strip()
            children = device.get("children") or []
            uuid = str(device.get("uuid") or "").strip()
            fs = str(device.get("fstype") or "")
            # A parent disk with partitions/RAID/LVM is represented by its usable leaves.
            if not children and kind not in {"loop", "rom"} and fs not in {"swap", "linux_raid_member", "LVM2_member", "crypto_LUKS"}:
                partition = re.search(r"(\d+)$", str(device.get("kname") or device.get("name") or ""))
                suffix = ":" + partition[1] if kind == "part" and partition else ""
                source = "uuid:" + uuid if uuid else "serial:" + serial + suffix if serial else "device:" + str(device.get("name"))
                identifier = hashlib.sha256(source.encode()).hexdigest()[:32]
                if identifier not in seen:
                    seen.add(identifier)
                    candidates = mounts.get(device.get("maj:min"), []) or sources.get(str(Path(device["name"]).resolve()), [])
                    mount = sorted(candidates)[0][2] if candidates else ""
                    item = dict(id=identifier, device=str(device.get("name") or ""),
                                label=str(device.get("label") or ""), uuid=uuid,
                                stable_identity=bool(uuid or serial), model=str(device.get("model") or "").strip(),
                                mountpoint=mount, filesystem=fs, state="mounted" if mount else "unmounted",
                                capacity_bytes=int(device.get("size") or 0), used_bytes=None,
                                free_bytes=None, used_percent=None)
                    if mount:
                        try:
                            usage = self.usage(mount)
                            if usage.total <= 0:
                                raise ValueError("invalid filesystem capacity")
                            item.update(capacity_bytes=usage.total, used_bytes=usage.used,
                                        free_bytes=usage.free, used_percent=round(100 * usage.used / usage.total, 1))
                        except (OSError, ValueError):
                            item["state"] = "unknown"
                    result.append(item)
            for child in children:
                walk(child, serial, str(device.get("name") or parent))

        for device in devices:
            walk(device)
        return sorted(result, key=lambda d: (d["mountpoint"] != "/", d["mountpoint"], d["device"]))



class StorageCache:
    def __init__(self, path, reader=disk_usage, clock=time.time):
        self.path, self.reader, self.clock = Path(path), reader, clock
        self.values = None
        self.observed_at = None
        self.next_update_at = 0
        self.error = False
        try:
            data = json.loads(self.path.read_text())
            if not isinstance(data,dict):raise ValueError('invalid cache')
            retry = data.get('next_update_at')
            if data.get('error') is True and isinstance(retry,(int,float)) and self.clock() < retry <= self.clock()+3600:
                self.error, self.next_update_at = True, retry
            values = data['values']
            if all(isinstance(values[k], int) and not isinstance(values[k], bool) and values[k] >= 0 for k in ('total', 'used', 'free')) and values['total'] > 0 and values['used'] <= values['total']:
                at = data['observed_at']
                if isinstance(at, (int, float)) and 0 <= self.clock() - at < 90000:
                    self.values, self.observed_at = values, at
                    if not self.error:self.next_update_at = at + 604800
        except (OSError, ValueError, KeyError, TypeError):
            pass

    def sample(self):
        now = self.clock()
        if now < self.next_update_at:
            return self.values
        try:
            usage = self.reader('/')
            if usage.total <= 0 or not 0 <= usage.used <= usage.total or usage.free < 0:
                raise ValueError('invalid capacity')
        except (OSError, ValueError):
            self.error = True
            self.next_update_at = now + 3600
            self.persist()
            return self.values
        self.values = dict(total=usage.total, used=usage.used, free=usage.free)
        self.observed_at = now
        self.next_update_at = now + 604800
        self.error = False
        self.persist()
        return self.values

    def persist(self):
        try:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            self.path.write_text(json.dumps(dict(values=self.values, observed_at=self.observed_at,
                                                next_update_at=self.next_update_at, error=self.error)))
        except OSError:
            # The in-memory sample remains valid even if the RAM cache cannot be persisted.
            pass