#!/usr/bin/env python3
"""merge_rootfs.py -- assemble a release root filesystem tarball from a JSON spec.

  python merge_rootfs.py <spec.json> <out.tar.gz> [--verify]

Works on Windows or Linux without root: tar members are copied as TarInfo, so
owners, modes, symlinks, hard links and device nodes survive exactly (an
extract-and-retar on Windows, or as non-root on Linux, would lose them).
Later layers override earlier ones by path; the output is deterministic
(fixed mtime for added entries, gzip header mtime 0), so the same spec gives
the same archive bytes.

spec.json (paths relative to the spec file unless absolute):
{
  "mtime": 1792022400,                       # added entries; optional
  "layers": [ "rootfs.tar.gz", "overlay.tar.gz" ],   # applied in order
  "remove": [ "etc/dropbear/dropbear_rsa_host_key" ], # paths (and their subtrees)
  "files":  [ {"path": "usr/bin/app", "src": "bin/app", "mode": "0755"} ],
  "links":  [ {"path": "etc/app/current.cfg", "target": "/etc/app/a.cfg"} ],
  "dirs":   [ "aw/dpd_curve" ],
  "expect": { "usr/bin/app": "<sha256 or a prefix>" }   # checked with --verify
}
"""
import gzip, hashlib, io, json, os, sys, tarfile


def norm(n):
    while n.startswith("./"):
        n = n[2:]
    n = n.strip("/")
    return "./" + n if n not in ("", ".") else "./"


class Tree:
    def __init__(self):
        self.m, self.order = {}, []

    def put(self, ti, data):
        ti.name = norm(ti.name)
        if ti.islnk():
            ti.linkname = norm(ti.linkname)
        if ti.name not in self.m:
            self.order.append(ti.name)
        self.m[ti.name] = (ti, data)

    def remove(self, p):
        p = norm(p)
        for n in list(self.m):
            if n == p or n.startswith(p + "/"):
                del self.m[n]
        self.order = [n for n in self.order if n in self.m]


def info(name, kind, mode, mtime):
    ti = tarfile.TarInfo(name); ti.type = kind; ti.mode = mode; ti.mtime = mtime
    ti.uid = ti.gid = 0; ti.uname = ti.gname = "root"
    return ti


def main(argv):
    if len(argv) < 3:
        sys.exit(__doc__)
    spec_path, out = argv[1], argv[2]
    base = os.path.dirname(os.path.abspath(spec_path))
    rel = lambda p: p if os.path.isabs(p) else os.path.join(base, p)
    spec = json.load(open(spec_path))
    mt = int(spec.get("mtime", 1700000000))
    t = Tree()
    for layer in spec.get("layers", []):
        with tarfile.open(rel(layer), "r:*") as tf:
            for ti in tf.getmembers():
                t.put(ti, tf.extractfile(ti).read() if ti.isfile() else None)
    for p in spec.get("remove", []):
        t.remove(p)
    for f in spec.get("files", []):
        d = open(rel(f["src"]), "rb").read()
        ti = info(f["path"], tarfile.REGTYPE, int(f.get("mode", "0644"), 8), mt); ti.size = len(d)
        t.put(ti, d)
    for l in spec.get("links", []):
        ti = info(l["path"], tarfile.SYMTYPE, 0o777, mt); ti.linkname = l["target"]
        t.put(ti, None)
    for d in spec.get("dirs", []):
        parts = norm(d)[2:].split("/")
        for k in range(1, len(parts) + 1):          # parents first, keep existing ones
            p = norm("/".join(parts[:k]))
            if p not in t.m:
                t.put(info(p, tarfile.DIRTYPE, 0o755, mt), None)
    raw = io.BytesIO()
    with tarfile.open(fileobj=raw, mode="w", format=tarfile.GNU_FORMAT) as tf:
        for n in t.order:
            ti, data = t.m[n]
            tf.addfile(ti, io.BytesIO(data) if data is not None else None)
    with open(out, "wb") as fo, gzip.GzipFile(filename="", mode="wb", fileobj=fo, mtime=0) as gz:
        gz.write(raw.getvalue())
    n_dev = sum(1 for n in t.order if t.m[n][0].isdev())
    print(f"{len(t.order)} members ({n_dev} device nodes) -> {out}")
    print("sha256", hashlib.sha256(open(out, "rb").read()).hexdigest())
    if "--verify" in argv:
        bad = 0
        for p, want in spec.get("expect", {}).items():
            e = t.m.get(norm(p))
            got = hashlib.sha256(e[1]).hexdigest() if e and e[1] is not None else None
            ok = got is not None and got.startswith(want.lower())
            bad += not ok
            print(("OK  " if ok else "BAD ") + (got or "missing")[:16], p)
        for l in spec.get("links", []):
            e = t.m.get(norm(l["path"]))
            ok = e is not None and e[0].issym() and e[0].linkname == l["target"]
            bad += not ok
            print(("OK  " if ok else "BAD ") + "link", l["path"], "->", l["target"])
        if bad:
            sys.exit(f"{bad} check(s) failed")


if __name__ == "__main__":
    main(sys.argv)