import struct
path = input("Path to 00000001.app: ")
raw = open(path, "rb").read()
print(f"file size {len(raw):#x}")
exefs_off=struct.unpack("<I",raw[0x1a0:0x1a4])[0] * 0x200
romfs_off=struct.unpack("<I",raw[0x1b0:0x1b4])[0] * 0x200
romfs_size=struct.unpack("<I",raw[0x1b4:0x1b8])[0] * 0x200
comp=raw[0x20d] & 1
print(f"ExeFS @{exefs_off:#x}  RomFS off={romfs_off:#x} size={romfs_size:#x}  code-compressed={comp}")
print("ExeFS files:")
for k in range(10):
    e=exefs_off + k * 16
    name=raw[e:e + 8].rstrip(b'\0')
    if not name:
        continue

    fo = struct.unpack("<I",raw[e + 8:e + 12])[0]
    fs = struct.unpack("<I",raw[e + 12:e + 16])[0]

    print(f"  {name.decode(errors='replace'):8} off={fo:#x} size={fs:#x} file@{exefs_off + 0x200 + fo:#x}")

def decompress(data):
    data = bytearray(data)
    dlen = len(data)
    btb = struct.unpack_from('<I',data, dlen - 8)[0]
    add = struct.unpack_from('<I',data, dlen - 4)[0]
    dec = dlen + add;
    res = bytearray(data) + bytearray(add)
    c = btb & 0xffffff
    hdr = btb >> 24
    stop = dlen - c
    ip = dlen - hdr
    op = dec

    while ip>stop:
        ip-=1; ctrl=res[ip]
        for _ in range(8):
            if ctrl&0x80:
                ip-=1; hi=res[ip]; ip-=1; lo=res[ip]; n=(hi>>4)+3; disp=(((hi&0xf)<<8)|lo)+3
                for _ in range(n): op-=1; res[op]=res[op+disp]
            else: ip-=1; op-=1; res[op]=res[ip]
            ctrl=(ctrl<<1)&0xff
            if ip<=stop: break
    return bytes(res)

# .code
ce=exefs_off; crel=struct.unpack("<I",raw[ce+8:ce+12])[0]; csz=struct.unpack("<I",raw[ce+12:ce+16])[0]
code=raw[exefs_off+0x200+crel:exefs_off+0x200+crel+csz]
coded=decompress(code) if comp else code
open("nwm_safe.code","wb").write(coded)
print(f"\n.code {len(code):#x} -> decompressed {len(coded):#x}")

targets={"stubdata 0x524C00":0x524C00,"stubcode 0x527000":0x527000,"stubcode 0x515000":0x515000,
         "main 0x502400":0x502400,"database 0x53FE18":0x53FE18,"database 0x52D944":0x52D944,
         "romver 0x2300006f":0x2300006f,"eeprom 0x609C0202":0x609C0202,"flags 0x80000001":0x80000001}
for label,blob in [("RAW .app",raw),("decompressed .code",coded)]:
    print(f"\n== {label} ==")
    for name,val in targets.items():
        offs=[i for i in range(0,len(blob)-4) if blob[i:i+4]==struct.pack("<I",val)]
        if offs: print(f"  {name:22}: {len(offs)} @ {', '.join(hex(o) for o in offs[:8])}")

code = open("nwm_safe.code", "rb").read()
stubdata = code[0x4f698:0x4f698+0x38]
stubcode = code[0x4f380:0x4f380+0x316]
main1    = code[0x3c878:0x3c878+0xfd3]
database = code[0x4f198:0x4f198+0x1e8]
HI, DB_DST = 0x520000, 0x53fe18
WRITE, LZ, EXEC = 0, 1, 2
parts = [
    (WRITE, 0x524c00, stubdata),                   # stub/data
    (WRITE, 0x527000, stubcode),                   # stub/code
    (EXEC,  0x927000, b""),                        # execute stub (0x527000 + 0x400000)
    (LZ,    0x524c00, main1),                      # LZ main
    (WRITE, 0x524c00, stubdata),                   # wifiboot: re-upload stub/data after LZ ("again")
    (WRITE, DB_DST,   database),                   # database
    (WRITE, HI+0x18,  struct.pack("<I", DB_DST)),  # host_interest[0x18] = &database
]

hdr  = b"AR6014FW" + struct.pack("<II", 1, len(parts))
desc = b"".join(struct.pack("<III", m, a, len(d)) for m,a,d in parts)
data = b"".join(d for _, _, d in parts)
open("fw-ar6014.bin", "wb").write(hdr + desc + data)
print(f"container: {len(hdr + desc + data)} bytes, {len(parts)} parts")
for m,a,d in parts:
    print(f"  {['WRITE','LZ','EXEC'][m]:5} 0x{a:08x} len=0x{len(d):04x}")