storage/diskimage.py

"""
Whole-drive imaging into a self-healing .pt.

Scans a block device (Windows: r"\\\\.\\C:" or r"\\\\.\\PhysicalDrive0"; may need
Administrator) or an image file, and turns it into a self-healing .pt archive:
content-defined chunks, Reed-Solomon shards, per-shard SHA-256 (Step 1-2 +
archive.py). It self-heals bit-rot in the archive and reconstructs the image
bit-exact.

Streaming by segment keeps memory bounded, so it scales past RAM:
  * small input (<= one segment)  -> a single self-contained  <out>.pt
  * large drive                    -> <out>.pt manifest + <out>.pt.part0, part1, ...

HONEST LIMITS (read these):
  * RS adds redundancy: the .pt is ~ (k+m)/k x the imaged data (e.g. 1.5x). This
    does NOT shrink a drive -- you cannot beat entropy.
  * Imaging an "entire HDD" means writing >= 1.5x its used size to disk and takes
    real time; use --part-mb and expect a big output.
  * Raw whole-disk reads need Administrator and the disk should be idle/unmounted
    (image a partition, or a disk you are not writing to).
  * RS heals the ARCHIVE against future bit-rot. Sectors unreadable at scan time
    are recorded as bad, not invented.
"""
from __future__ import annotations
import os
import torch
from . import archive, scan

DEFAULT_PART = 64 * 1024 * 1024        # 64 MB segments


def _read_segment(f, part_bytes: int, sector: int):
    """Read up to part_bytes in sector-sized reads. -> (data, bad_sectors, eof)."""
    data = bytearray()
    bad = []
    idx = 0
    while len(data) < part_bytes:
        try:
            blk = f.read(sector)
        except OSError:
            bad.append(idx)
            data += bytes(sector)
            idx += 1
            try:
                f.seek(sector, os.SEEK_CUR)
            except OSError:
                return bytes(data), bad, True
            continue
        if not blk:
            return bytes(data), bad, True
        data += blk
        idx += 1
        if len(blk) < sector:
            return bytes(data), bad, True
    return bytes(data), bad, False


def create(device: str, out_pt: str, k: int = 4, m: int = 2,
           chunk: int = archive.CHUNK, part_mb: int = 64,
           sector: int = scan.SECTOR, label: str | None = None) -> dict:
    label = label or os.path.basename(device.rstrip("\\/")) or "disk"
    part_bytes = part_mb * 1024 * 1024
    with open(device, "rb", buffering=0) as f:
        seg, bad, eof = _read_segment(f, part_bytes, sector)
        if eof:                                    # whole device fit in one segment
            arc = archive.build(seg, label=label, k=k, m=m, bad_sectors=bad)
            arc["meta"]["device"] = device
            torch.save(arc, out_pt)
            return {"mode": "single", "bytes": len(seg), "bad_sectors": len(bad),
                    "parts": 1, "out": out_pt}
        parts, n, total_bad = [], 0, 0
        while True:
            p = f"{out_pt}.part{n}"
            arc = archive.build(seg, label=f"{label}.part{n}", k=k, m=m, bad_sectors=bad)
            torch.save(arc, p)
            parts.append({"path": os.path.basename(p), "size": len(seg)})
            total_bad += len(bad)
            n += 1
            if eof:
                break
            seg, bad, eof = _read_segment(f, part_bytes, sector)
            if not seg and eof:
                break
    manifest = {"format": "neural-storage/diskimage-1", "device": device, "label": label,
                "k": k, "m": m, "chunk": chunk, "parts": parts,
                "total": sum(p["size"] for p in parts), "bad_sectors": total_bad}
    torch.save(manifest, out_pt)
    return {"mode": "multi", "bytes": manifest["total"], "bad_sectors": total_bad,
            "parts": len(parts), "out": out_pt}


def _is_manifest(obj) -> bool:
    return isinstance(obj, dict) and obj.get("format", "").startswith("neural-storage/diskimage")


def verify(out_pt: str) -> dict:
    obj = archive.load(out_pt)
    if not _is_manifest(obj):
        return archive.verify(obj)
    base = os.path.dirname(out_pt)
    agg = {"chunks": 0, "corrupted_shards": 0, "repairable_chunks": 0, "lost_chunks": 0}
    for p in obj["parts"]:
        h = archive.verify(archive.load(os.path.join(base, p["path"])))
        for kk in agg:
            agg[kk] += h[kk]
    agg["healthy"] = agg["corrupted_shards"] == 0 and agg["lost_chunks"] == 0
    return agg


def heal(out_pt: str) -> int:
    obj = archive.load(out_pt)
    if not _is_manifest(obj):
        n = archive.heal(obj)
        torch.save(obj, out_pt)
        return n
    base = os.path.dirname(out_pt)
    total = 0
    for p in obj["parts"]:
        pth = os.path.join(base, p["path"])
        arc = archive.load(pth)
        n = archive.heal(arc)
        if n:
            torch.save(arc, pth)
        total += n
    return total


def restore(out_pt: str, out_file: str) -> int:
    obj = archive.load(out_pt)
    with open(out_file, "wb") as w:
        if not _is_manifest(obj):
            data = archive.restore(obj)
            w.write(data)
            return len(data)
        base = os.path.dirname(out_pt)
        total = 0
        for p in obj["parts"]:
            data = archive.restore(archive.load(os.path.join(base, p["path"])))
            w.write(data)
            total += len(data)
        return total