#!/usr/bin/env python3
"""ffu_extract.py — извлечение разделов из FFU v1 payload в файлы/eMMC-образ.
Модель (проверена на RM1127 orig):
- payload data chunk k лежит в файле по смещению DATA0 + k*131072,
  DATA0 = первая граница чанка после заголовков (здесь 0xE0000).
- WriteDescriptor entries по порядку; entries с method-0 потребляют
  чанки данных (holes method-2 пропускаются); disk chunk = LBA//256.
Использование:
  ./tools/ffu_extract.py RM1127.ffu out-dir/ [LBA0 LBA1 ...]
  без LBA — извлекает SBL1/DBI/UEFI/RPM/QSEE/QHEE/WINSECAPP/TZAPPS/PLAT/EFIESP
  ./tools/ffu_extract.py RM1127.ffu --emmc-user out-user.bin
"""
import os
import struct
import sys

CHUNK = 131072


def load(ffu):
    import mmap
    raw = open(ffu, "rb")
    f = mmap.mmap(raw.fileno(), 0, access=mmap.ACCESS_READ)
    sec = f.read(32)
    _, _, chunk_kb, _, cat, hsh = struct.unpack("<I12sIIII", sec)
    assert chunk_kb * 1024 == CHUNK
    # store header: ASCII PlatformId inside it; header starts 12B before.
    # (manifest has one too — pick the occurrence with sane blkSize)
    st = -1
    i = 0
    while True:
        i = f.find(b"Microsoft.MSM", i)
        if i < 0:
            break
        vals = struct.unpack("<11I", f[i - 12 + 204:i - 12 + 248])
        if vals[0] == CHUNK:
            st = i - 12
            break
        i += 1
    assert st > 0, "store header not found"
    wcnt = vals[1]
    pos = st + 248
    disk = {}   # diskchunk -> datachunk
    dc = 0
    for _ in range(wcnt):
        a, _ = struct.unpack("<II", f[pos:pos + 8])
        locs = [struct.unpack("<II", f[pos + 8 + k * 8:pos + 16 + k * 8])
                for k in range(a)]
        pos += 8 + 8 * a
        for m, idx in locs:
            if m == 0:
                disk[idx] = dc
                dc += 1
    data0 = st + 248 + vals[2]
    data0 += (CHUNK - data0 % CHUNK) % CHUNK
    return f, data0, disk


# имя -> (lba0, секторы) из tools/ffu_parse.py по RM1127 orig
PARTS = {
    "SBL1": (90112, 1024), "DBI": (94208, 64), "UEFI": (98304, 4096),
    "RPM": (102400, 1024), "QSEE": (106496, 2048), "QHEE": (110592, 1024),
    "WINSECAPP": (114688, 1024), "TZAPPS": (118784, 32768),
    "PLAT": (212992, 16384), "EFIESP": (229376, 131072),
    "DPP": (4096, 16384), "MODEM_FSG": (20480, 6144),
    "SSD": (73728, 32), "UEFI_BS_NV": (77824, 512),
    "UEFI_RT_NV": (81920, 512), "BOOTMODE": (86016, 32),
    "SEC": (372736, 32),
}


def chunk_of(f, data0, disk, dchunk):
    if dchunk not in disk:
        return None
    f.seek(data0 + disk[dchunk] * CHUNK)
    return f.read(CHUNK)


def main():
    ffu, out = sys.argv[1], sys.argv[2]
    f, data0, disk = load(ffu)
    print("data0:", hex(data0), "mapped chunks:", len(disk))
    if out == "--emmc-user":
        dest = sys.argv[3]
        with open(dest, "wb") as o:
            o.truncate(15154944 * 512)
            for name, (lba, n) in PARTS.items():
                for s in range(0, n, 256):
                    c = chunk_of(f, data0, disk, (lba + s) // 256)
                    if c is None:
                        continue
                    o.seek((lba + s) * 512)
                    o.write(c[:min(n - s, 256) * 512])
        print("wrote", dest)
        return
    os.makedirs(out, exist_ok=True)
    names = sys.argv[3:] or list(PARTS)
    for name in names:
        lba, n = PARTS[name]
        with open(os.path.join(out, name + ".bin"), "wb") as o:
            for s in range(0, n, 256):
                c = chunk_of(f, data0, disk, (lba + s) // 256)
                o.write(c[:min(n - s, 256) * 512] if c else b"\x00" * 512)
        print("wrote", name)


if __name__ == "__main__":
    main()