File size: 5,286 Bytes
bdfb884
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
"""
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