#!/usr/bin/env python3
"""Archive rootless staging with package ownership and portable Linux capabilities."""
import json
import os
import pathlib
import shutil
import stat
import struct
import subprocess
import sys
import tarfile

from rootfs_lib import digest, package_db, rooted

root, work, *caches = (pathlib.Path(arg).resolve() for arg in sys.argv[1:])
db = package_db(root)
inputs = work / "packages"
inputs.mkdir()
owners = {}
records = []
scripts = []
for name, props in sorted(db.items()):
    filename = f'{props["pkgver"]}.{props["architecture"]}.xbps'
    candidates = [cache / filename for cache in caches if (cache / filename).is_file()]
    if not candidates:
        sys.exit(f"ERROR: selected package input is missing: {filename}")
    source = candidates[0]
    checksum = digest(source)
    expected = props.get("filename-sha256")
    if expected and expected != checksum:
        sys.exit(f"ERROR: package input changed after installation: {filename}")
    shutil.copyfile(source, inputs / filename)
    if source.with_suffix(".xbps.sig2").is_file():
        shutil.copyfile(source.with_suffix(".xbps.sig2"), inputs / (filename + ".sig2"))
    records.append({"name": name, "version": props["pkgver"], "architecture": props["architecture"], "sha256": checksum, "filename": filename})
    # Python 3.14 understands XBPS's zstd archives without third-party modules.
    with tarfile.open(source, "r:*") as archive:
        for member in archive:
            path = member.name.removeprefix("./")
            if path in ("INSTALL", "REMOVE"):
                scripts.append(f"### {filename}: {path}\n" + archive.extractfile(member).read().decode())
            if path in ("INSTALL", "REMOVE", "props.plist", "files.plist") or not path:
                continue
            if path.startswith("/") or ".." in pathlib.PurePosixPath(path).parts:
                sys.exit(f"ERROR: unsafe package path: {path}")
            # Resolve parent aliases (/bin -> /usr/bin), not a symlink leaf.
            parent = rooted(root, str(pathlib.PurePosixPath(path).parent))
            canonical = str((parent / pathlib.PurePosixPath(path).name).relative_to(root))
            value = (member.uid, member.gid)
            if canonical in owners and owners[canonical] != value:
                sys.exit(f"ERROR: conflicting package ownership: {canonical}")
            owners[canonical] = value
# util-linux's INSTALL assigns tty ownership after extraction. The single-user
# namespace cannot represent every target group; restore this tested policy.
for name in ("usr/bin/wall", "usr/bin/write"):
    if (root / name).exists():
        owners[name] = (0, 5)
# libutempter's INSTALL assigns utmp after extraction. A single-UID build
# namespace cannot apply that group; exporting SGID root would be incorrect.
if (root / "usr/lib/utempter/utempter").exists():
    groups = {line.split(":")[0]: int(line.split(":")[2]) for line in (root / "etc/group").read_text().splitlines() if line and not line.startswith("#")}
    owners["usr/lib/utempter/utempter"] = (0, groups["utmp"])
# XBPS restricts its setuid chroot helper to the dedicated xbuilder group.
owners["usr/bin/xbps-uchroot"] = (0, 101)
epoch = int(subprocess.check_output(["git", "-C", "vendor/void-packages", "show", "-s", "--format=%ct", "HEAD"]))
metadata = {}


def normalize(info):
    info.uid, info.gid = owners.get(info.name, (0, 0))
    info.uname = info.gname = ""
    info.mtime = epoch
    info.mode = stat.S_IMODE(info.mode)
    init_fifo = info.isfifo() and info.name.startswith("etc/s6-linux-init/current/run-image/")
    if not (info.isfile() or info.isdir() or info.issym() or info.islnk() or init_fifo):
        raise ValueError(f"Unexpected special file in image: {info.name}")
    path = root / info.name
    if not info.issym() and "security.capability" in os.listxattr(path):
        capability = os.getxattr(path, "security.capability")
        revision = struct.unpack_from("<I", capability)[0]
        # User namespaces store v3 capabilities with the builder's host UID.
        # An installed root filesystem needs the equivalent root-owned v2 form.
        if revision >> 24 == 3:
            capability = struct.pack("<I", (revision & 0xFFFFFF) | 0x02000000) + capability[4:20]
        if len(capability) != 20 or struct.unpack_from("<I", capability)[0] >> 24 != 2:
            raise ValueError(f"Unsupported capability encoding: {info.name}")
        info.pax_headers["SCHILY.xattr.security.capability"] = capability.decode("utf-8", "surrogateescape")
    metadata[info.name] = {"uid": info.uid, "gid": info.gid, "mode": oct(info.mode), "capability": info.pax_headers.get("SCHILY.xattr.security.capability", "").encode("utf-8", "surrogateescape").hex()}
    return info


output = work / "rootfs-aarch64.tar"
with tarfile.open(output, "w", format=tarfile.PAX_FORMAT) as archive:
    # EROFS tar import otherwise synthesizes / using the host UID and 0777.
    # Record the filesystem root explicitly, not only its children.
    root_info = tarfile.TarInfo('.')
    root_info.type = tarfile.DIRTYPE
    root_info.mode = 0o755
    archive.addfile(normalize(root_info))
    for path in sorted(root.iterdir()):
        archive.add(path, arcname=path.name, filter=normalize)
(work / "packages.json").write_text(json.dumps(records, indent=2) + "\n")
(work / "archive-metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
(work / "package-scripts.txt").write_text("\n".join(scripts))
(work / "packages.sha256").write_text("".join(f'{item["sha256"]}  packages/{item["filename"]}\n' for item in records))
(work / "rootfs.sha256").write_text(f"{digest(output)}  rootfs-aarch64.tar\n")
print(f"PASS: archive with restored ownership and capabilities ({output.stat().st_size} bytes)")