#!/usr/bin/env python3
"""递归恢复 exFAT 已删除目录树（只读镜像 → 输出到本地目录）。
用法: exfat_tree.py <img> <start_clu> <start_datalen> <outdir|--dry>
连续分配(NoFatChain=True)按物理连续读；否则尝试 FAT 链，FAT 已清则退回连续读。
"""
import struct, sys, os, math, json

IMG = sys.argv[1]
START_CLU = int(sys.argv[2])
START_LEN = int(sys.argv[3])
OUT = sys.argv[4]
DRY = (OUT == '--dry')

f = open(IMG, 'rb')
def rd(off, n):
    f.seek(off); return f.read(n)

bs_off = next(off for off in range(0, 1 << 20, 512)
              if rd(off + 3, 8) == b'EXFAT   ')
bs = rd(bs_off, 512)
bps = 1 << bs[0x6C]
spc = 1 << bs[0x6D]
fat_off = struct.unpack_from('<I', bs, 0x50)[0]
heap_off = struct.unpack_from('<I', bs, 0x58)[0]
clu_count = struct.unpack_from('<I', bs, 0x5C)[0]
FAT_BASE = bs_off + fat_off * bps
HEAP_BASE = bs_off + heap_off * bps
CLU = bps * spc
MAX_CLU = clu_count + 1

def clu_off(n): return HEAP_BASE + (n - 2) * CLU
def read_clu(n): return rd(clu_off(n), CLU)
def fat_entry(n):
    f.seek(FAT_BASE + 4 * n); return struct.unpack('<I', f.read(4))[0]

def read_contig(start, length):
    if start < 2 or start > MAX_CLU:
        return None
    nclu = max(1, math.ceil(length / CLU)) if length else 1
    if start + nclu - 1 > MAX_CLU:
        return None
    data = b''.join(read_clu(start + i) for i in range(nclu))
    return data[:length] if length else data

def read_chain_or_contig(start, length, nofat):
    """先按 FAT 链读（若可用且长度吻合），否则连续读"""
    if not nofat:
        clus, c, seen, got = [], start, set(), 0
        while 2 <= c < 0xFFFFFFF8 and c not in seen and len(clus) < 1 << 20:
            seen.add(c); clus.append(c); got += CLU
            if got >= length: break
            c = fat_entry(c)
        if clus and got >= length:
            data = b''.join(read_clu(x) for x in clus)
            return data[:length]
    return read_contig(start, length)

def parse_entries(raw):
    ents = []
    for i in range(0, len(raw) - 31, 32):
        e = raw[i:i + 32]
        t = e[0]
        if t == 0x00:
            break
        ents.append({'t': t, 'code': t & 0x1F, 'cat2': bool(t & 0x40),
                     'live': bool(t & 0x80), 'raw': e})
    return ents

def walk(ents):
    out, i = [], 0
    while i < len(ents):
        e = ents[i]
        if not e['cat2'] and e['code'] == 5:
            sec = e['raw'][1]
            attrs = struct.unpack_from('<H', e['raw'], 4)[0]
            name, stream = '', None
            for j in range(i + 1, min(i + 1 + sec, len(ents))):
                s = ents[j]
                if s['cat2'] and s['code'] == 0:
                    stream = s['raw']
                elif s['cat2'] and s['code'] == 1:
                    name += s['raw'][2:32].decode('utf-16-le', 'replace')
            d = {'name': name, 'attrs': attrs, 'isdir': bool(attrs & 0x10),
                 'live': e['live'], 'clu': None, 'size': 0, 'nofat': None}
            if stream is not None:
                d['nofat'] = bool(stream[1] & 2)
                d['namelen'] = stream[3]
                d['clu'] = struct.unpack_from('<I', stream, 0x14)[0]
                d['size'] = struct.unpack_from('<Q', stream, 0x18)[0]
                d['name'] = d['name'][:d['namelen']]
            out.append(d)
            i += 1 + sec
        else:
            i += 1
    return out

def list_dir(cluster, datalen):
    raw = read_contig(cluster, datalen) or b''
    return walk(parse_entries(raw))

manifest, total_files, total_bytes, visited = [], 0, 0, set()

def rec(cluster, datalen, relpath, depth=0):
    global total_files, total_bytes
    if depth > 32 or cluster in visited:
        return
    visited.add(cluster)
    items = list_dir(cluster, datalen)
    if not DRY:
        os.makedirs(os.path.join(OUT, relpath), exist_ok=True)
    for d in items:
        name = d['name'].replace('\x00', '').strip()
        if not name or name in ('.', '..'):
            continue
        sub = os.path.join(relpath, name)
        if d['isdir']:
            print(f"{'  '*depth}DIR  {sub}/  (首簇={d['clu']} 大小={d['size']})")
            rec(d['clu'], d['size'], sub, depth + 1)
        else:
            data = read_chain_or_contig(d['clu'], d['size'], d['nofat']) if d['size'] else b''
            ok = data is not None and len(data) == d['size']
            if not DRY and ok:
                p = os.path.join(OUT, sub)
                os.makedirs(os.path.dirname(p), exist_ok=True)
                open(p, 'wb').write(data)
            total_files += 1
            total_bytes += d['size']
            manifest.append({'path': sub, 'size': d['size'], 'cluster': d['clu'],
                             'nofat': d['nofat'], 'read_ok': ok})
            flag = '' if ok else '  <<读失败/越界>>'
            print(f"{'  '*depth}FILE {sub}  {d['size']}B 首簇={d['clu']}{flag}")

rec(START_CLU, START_LEN, 'zyw')
print(f'\n合计: {total_files} 个文件, {total_bytes} 字节 '
      f'({total_bytes/1048576:.1f} MiB)')
outjson = '/persistent/usb-recover/manifest.json'
try:
    json.dump({'files': manifest, 'total_files': total_files,
               'total_bytes': total_bytes}, open(outjson, 'w'),
              ensure_ascii=False, indent=1)
    print(f'清单已写: {outjson}')
except Exception as e:
    print('清单写入失败:', e)