#!/usr/bin/env python3
"""Construit une image NAND pour l'émulateur iPod Touch 2G avec nos apps / paquets dedans.

  ipod-mkimage [-b BASE] [-o NOM] fichier.app|fichier.ipa|fichier.deb|dossier/ [...]

  -b BASE  firmware de base (défaut : 5H11a = iPhone OS 2.2.1 ; 5F138 = 2.1.1) — utilise ~/ipod-emu/images/BASE/rootfs.img
  -o NOM   nom de l'image produite (défaut : BASE-apps) — ensuite : ipod-emu start -f -w NOM

Les .app et .ipa sont installés dans /Applications (apps système : pas besoin d'installd,
la signature ldid suffit), les .deb sont décompressés à la racine comme dpkg le ferait, un dossier est copié tel quel
à la racine (arborescence « overlay » : ses binaires Mach-O sont signés avec ldid).
"""
import argparse, io, lzma, gzip, bz2, os, re, shutil, subprocess, sys, tarfile, tempfile, zipfile

EMU = os.environ.get("IPOD_EMU_HOME", os.path.expanduser("~/ipod-emu"))
HFSPLUS = os.path.join(EMU, "xpwn/b/hfs/hfsplus")
MKNAND = os.path.join(EMU, "gen-v1/mknand")
MACHO = (b"\xce\xfa\xed\xfe", b"\xca\xfe\xba\xbe")


def hfs(img, *args, check=True):
    r = subprocess.run([HFSPLUS, img, *args], capture_output=True, text=True)
    if check and r.returncode:
        sys.exit(f"hfsplus {' '.join(args)} : {r.stderr or r.stdout}")
    return r.stdout


def exists(img, path):
    return "No such file" not in hfs(img, "ls", path, check=False) and hfs(img, "ls", path, check=False).strip() != ""


def mkdirs(img, path):
    cur = ""
    for part in path.strip("/").split("/"):
        cur += "/" + part
        if not exists(img, cur):
            hfs(img, "mkdir", cur)


def add_tree(img, src, dest):
    """Copie récursive src -> dest dans l'image, en gardant les bits d'exécution."""
    mkdirs(img, dest)
    for root, dirs, files in os.walk(src):
        rel = os.path.relpath(root, src)
        tgt = dest if rel == "." else f"{dest}/{rel}"
        for d in dirs:
            p = os.path.join(root, d)
            if os.path.islink(p):
                continue
            hfs(img, "mkdir", f"{tgt}/{d}", check=False)
        for f in files:
            p = os.path.join(root, f)
            if os.path.islink(p):
                print(f"  (lien symbolique ignoré : {p})")
                continue
            hfs(img, "add", p, f"{tgt}/{f}")
            with open(p, "rb") as fh:
                macho = fh.read(4) in MACHO
            if macho or os.stat(p).st_mode & 0o111:
                hfs(img, "chmod", "755", f"{tgt}/{f}")


def sign(tree):
    for root, _, files in os.walk(tree):
        for f in files:
            p = os.path.join(root, f)
            if os.path.islink(p):
                continue
            with open(p, "rb") as fh:
                if fh.read(4) not in MACHO:
                    continue
            # re-signe en gardant les entitlements déjà présents (ldid -S seul les efface)
            ents = subprocess.run(["ldid", "-e", p], capture_output=True).stdout
            if ents.strip():
                with tempfile.NamedTemporaryFile(suffix=".plist", delete=False) as t:
                    t.write(ents)
                subprocess.run(["ldid", "-S" + t.name, p], check=True)
                os.unlink(t.name)
            else:
                subprocess.run(["ldid", "-S", p], check=True)


def deb_data(path, out):
    with open(path, "rb") as f:
        assert f.read(8) == b"!<arch>\n", f"{path} : pas un .deb"
        while hdr := f.read(60):
            name, size = hdr[:16].decode().strip().rstrip("/"), int(hdr[48:58])
            data = f.read(size)
            f.read(size % 2)
            if name.startswith("data.tar"):
                opener = {".gz": gzip.decompress, ".bz2": bz2.decompress, ".xz": lzma.decompress, ".lzma": lzma.decompress}
                ext = os.path.splitext(name)[1]
                raw = opener.get(ext, lambda x: x)(data)
                tarfile.open(fileobj=io.BytesIO(raw)).extractall(out, filter="fully_trusted")
                return
    sys.exit(f"{path} : pas de data.tar")


def main():
    p = argparse.ArgumentParser(usage=__doc__)
    p.add_argument("-b", "--base", default="5H11a")
    p.add_argument("-o", "--out")
    p.add_argument("items", nargs="+")
    a = p.parse_args()
    out = a.out or f"{a.base}-apps"
    base_dir = os.path.join(EMU, "images", a.base)
    base_img = os.path.join(base_dir, "rootfs.img")
    if not os.path.exists(base_img):
        sys.exit(f"Pas de {base_img}")
    out_dir = os.path.join(EMU, "images", out)
    os.makedirs(out_dir, exist_ok=True)
    for f in ("bootrom", "nor.bin", "nand-version"):
        src, dst = os.path.join(base_dir, f), os.path.join(out_dir, f)
        if os.path.exists(src) and not os.path.lexists(dst):
            os.symlink(os.path.realpath(src), dst)

    with tempfile.TemporaryDirectory() as tmp:
        img = os.path.join(out_dir, "rootfs.img")
        shutil.copyfile(base_img, img)
        for item in a.items:
            item = os.path.abspath(item.rstrip("/"))
            print(f"+ {os.path.basename(item)}")
            stage = tempfile.mkdtemp(dir=tmp)
            if item.endswith(".app") and os.path.isdir(item):
                shutil.copytree(item, os.path.join(stage, os.path.basename(item)), symlinks=True)
                sign(stage)
                add_tree(img, os.path.join(stage, os.path.basename(item)), f"/Applications/{os.path.basename(item)}")
            elif item.endswith(".ipa"):
                zipfile.ZipFile(item).extractall(stage)
                apps = [d for d in os.listdir(os.path.join(stage, "Payload")) if d.endswith(".app")]
                for app in apps:
                    src = os.path.join(stage, "Payload", app)
                    sign(src)
                    add_tree(img, src, f"/Applications/{app}")
            elif os.path.isdir(item):
                shutil.copytree(item, os.path.join(stage, "root"), symlinks=True)
                sign(os.path.join(stage, "root"))
                add_tree(img, os.path.join(stage, "root"), "/")
            elif item.endswith(".deb"):
                deb_data(item, stage)
                add_tree(img, stage, "/")
            else:
                sys.exit(f"Type non géré : {item}")
        print(hfs(img, "ls", "/Applications").splitlines()[-1])
        try:
            version = open(os.path.join(base_dir, "nand-version")).read().strip()
        except FileNotFoundError:
            version = "1"
        subprocess.run([MKNAND, img, os.path.join(out_dir, "nand"), version], check=True)
        opts_path = os.path.join(base_dir, "machine-opts")
        opts = open(opts_path).read().strip() if os.path.exists(opts_path) else ""
        opts = re.sub(r",nand-write=(?:on|off)", "", opts) + ",nand-write=on"
        out_opts = os.path.join(out_dir, "machine-opts")
        if os.path.islink(out_opts):
            os.unlink(out_opts)
        with open(out_opts, "w") as f:
            f.write(opts + "\n")
    print(f"Image prête : {out_dir}\n  ipod-emu start -f -w {out}")


if __name__ == "__main__":
    main()
