#!/usr/bin/env python3
"""exFAT 删除恢复工具（只读镜像）。
用法:
  exfat_tool.py <img> list <cluster> [datalen]
  exfat_tool.py <img> recover <cluster> <datalen> <outdir>
  exfat_tool.py <img> dumpfile <cluster> <datalen> <outfile>
删除后 FAT 链已清空，故按「物理连续」读取（exFAT 绝大多数为大块连续分配）。
"""
import struct, sys, os, math

IMG = sys.argv[1]
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]
FAT_BASE = bs_off + fat_off * bps
HEAP_BASE = bs_off + heap_off * bps
CLU = bps * spc
VOL_END = HEAP_BASE + struct.unpack_from('<I', bs, 0x5C)[0] * CLU

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_range(start_clu, length):
    """从 start_clu 起连续读取 length 字节"""
    nclu = max(1, math.ceil(length / CLU))
    data = b''.join(read_clu(start_clu + i) for i in range(nclu))
    return data[:length] if length else data

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'], 'sec': sec, 'clu': None, 'size': 0,
                 'valid': 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['valid'] = struct.unpack_from('<Q', stream, 0x08)[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_range(cluster, datalen)
    return walk(parse_entries(raw)), raw

cmd = sys.argv[2]

if cmd == 'list':
    cluster, datalen = int(sys.argv[3]), int(sys.argv[4])
    items, raw = list_dir(cluster, datalen)
    print(f'目录簇={cluster} 数据长度={datalen} ({datalen // CLU} 簇)')
    print(f'{len(items)} 个条目:')
    for d in items:
        tag = '在用  ' if d['live'] else '已删除'
        print(f"  [{tag}] {'DIR ' if d['isdir'] else 'FILE'} "
              f"name={d['name']!r} 首簇={d['clu']} 大小={d['size']} "
              f"有效={d['valid']} NoFatChain={d['nofat']} 属性=0x{d['attrs']:04x}")

elif cmd == 'dumpfile':
    cluster, datalen, outfile = int(sys.argv[3]), int(sys.argv[4]), sys.argv[5]
    data = read_range(cluster, datalen)
    open(outfile, 'wb').write(data)
    print(f'已导出 {len(data)} 字节 -> {outfile}')

elif cmd == 'recover':
    cluster, datalen, outdir = int(sys.argv[3]), int(sys.argv[4]), sys.argv[5]
    os.makedirs(outdir, exist_ok=True)
    items, raw = list_dir(cluster, datalen)
    n = 0
    for d in items:
        if d['isdir']:
            # 子目录：按 1 簇起递归（先只列一层）
            print(f"  [子目录] {d['name']} 首簇={d['clu']} 大小={d['size']}")
            continue
        if not d['clu'] or not d['size']:
            print(f"  [跳过] {d['name']} 无簇/无大小")
            continue
        data = read_range(d['clu'], d['size'])
        safe = d['name'].replace('/', '_').replace('\x00', '')
        path = os.path.join(outdir, safe)
        open(path, 'wb').write(data)
        n += 1
        print(f"  恢复 {safe}  {len(data)} 字节  首簇={d['clu']}")
    print(f'共恢复 {n} 个文件到 {outdir}')
else:
    print(__doc__)