#!/usr/bin/env python3
"""Fail closed on wrong ABI, forbidden init stacks, or unfinished package setup."""
import json
import os
import pathlib
import stat
import struct
import sys
import shlex
from fds_version import validate

from rootfs_lib import package_db, rooted


def audit(root, configured=True):
    db = package_db(root)
    required = "fds-base fds-base-files fds-init fds-dasungd fds-cli fds-cartridged fds-kernel fds-eink fds-dhcpcd xorg-server WindowMaker terminus-font xf86-input-libinput glibc glibc-locales coreutils binutils xbps bash s6 s6-rc s6-linux-init".split()
    for name in required:
        assert name in db, f"Missing required package: {name}"
    forbidden = ("busybox", "runit", "systemd", "musl", "base-system", "base-files")
    for name, props in db.items():
        assert not any(name == prefix or name.startswith(prefix + "-") for prefix in forbidden), f"Forbidden package: {name}"
        assert props["architecture"] in ("aarch64", "noarch"), f"Wrong package ABI: {name}"
        if configured:
            assert props["state"] == "installed", f"Unconfigured package: {name}"
    layout = json.loads((root / "usr/share/fds/layout.json").read_text())
    for name, target in layout["symlinks"].items():
        assert (root / name).is_symlink() and str((root / name).readlink()) == target, f"Wrong layout link: {name}"
    for name, mode in layout["directories"].items():
        path = root / name
        assert path.is_dir() and not path.is_symlink(), f"Missing directory: {name}"
        assert stat.S_IMODE(path.stat().st_mode) == int(mode, 8), f"Wrong directory mode: {name}"
    assert 'ID=fds\n' in (root / "usr/lib/os-release").read_text(), "Wrong OS identity"
    identity = dict(line.split('=', 1) for line in shlex.split((root / 'usr/lib/os-release').read_text(), comments=True))
    version = validate(identity['VERSION_ID'])
    assert identity['VERSION'] == version and identity['PRETTY_NAME'] == f'FDS/OS {version}', 'Inconsistent OS version'
    assert (root / 'etc/issue').read_text() == f'FDS/OS {version}\n', 'Inconsistent login version'
    for package in ('fds-base', 'fds-base-files', 'fds-init', 'fds-cli', 'fds-cartridged', 'fds-dasungd', 'fds-eink'):
        assert db[package]['pkgver'] == f'{package}-{version}_1', f'OS/package version mismatch: {package}'
    release = (root / "usr/share/fds/kernel-release").read_text().strip()
    assert (root / "usr/lib/modules" / release / "modules.dep").is_file(), "Missing matching kernel module index"
    assert (root / "etc/shadow").stat().st_mode & 0o777 == 0o600, "Shadow file is exposed"
    assert (root / "etc/shadow").read_text().startswith("root:!:"), "Default root account must be locked"
    for name in ("etc/sv", "etc/runit", "etc/systemd/system", "usr/bin/busybox", "usr/bin/runit", "usr/bin/runsv", "usr/bin/systemctl", "usr/lib/systemd/systemd"):
        path = root / name
        assert not path.exists() and not path.is_symlink(), f"Unexpected runtime: {name}"
    profile = (root / "usr/share/fds/image-profile").read_text().strip()
    assert profile in ("cli", "development", "recovery"), "Unknown image profile"
    for optional in ("desktop", "xserver", "desktop-session", "network", "dhcp"):
        assert not (root / f"etc/s6-rc/source/boot/contents.d/{optional}").exists(), "Optional service in boot bundle"
    if configured:
        for name in ("WindowMaker", "WMRootMenu", "WMWindowAttributes"):
            assert (root / "etc/WindowMaker" / name).read_bytes() == (root / "usr/share/fds/eink" / name).read_bytes(), "FDS grayscale defaults must be global"
    for name in ("fds-program", "fds-control"):
        assert (root / "usr/bin" / name).is_file(), f"Missing FDS interface: {name}"
    if profile == "development":
        assert "xorg-server-xvfb" in db and "xdotool" in db, "Missing development display diagnostics"
        for package in "gcc glibc-devel make cmake meson ninja pkg-config rust cargo git gdb strace vim".split():
            assert package in db, f"Missing development tool: {package}"
    assert (root / "etc/resolv.conf").is_symlink() and str((root / "etc/resolv.conf").readlink()) == "/run/fds/resolv.conf", "DNS must be volatile"
    elf_count = 0
    for directory, _, files in os.walk(root, followlinks=False):
        for name in files:
            path = pathlib.Path(directory) / name
            if path.is_symlink():
                continue
            if not path.is_file():
                assert stat.S_ISFIFO(path.stat().st_mode) and str(path.relative_to(root)).startswith("etc/s6-linux-init/current/run-image/"), f"Unexpected special file: {path}"
                continue
            with path.open("rb") as stream:
                header = stream.read(64)
                if header[:4] != b"\x7fELF":
                    continue
                elf_count += 1
                assert header[4:6] == b"\x02\x01" and struct.unpack_from("<H", header, 18)[0] == 183, f"Non-aarch64 ELF: {path.relative_to(root)}"
                offset = struct.unpack_from("<Q", header, 32)[0]
                size, count = struct.unpack_from("<HH", header, 54)
                for i in range(count):
                    stream.seek(offset + i * size)
                    ph = stream.read(size)
                    if struct.unpack_from("<I", ph)[0] == 3:
                        pos = struct.unpack_from("<Q", ph, 8)[0]
                        length = struct.unpack_from("<Q", ph, 32)[0]
                        stream.seek(pos)
                        interpreter = stream.read(length).rstrip(b"\0").decode()
                        assert interpreter in ("/lib/ld-linux-aarch64.so.1", "/lib64/ld-linux-aarch64.so.1", "/usr/lib/ld-linux-aarch64.so.1", "/usr/lib64/ld-linux-aarch64.so.1"), f"Wrong ELF interpreter: {path.relative_to(root)}: {interpreter}"
                        assert rooted(root, interpreter).is_file(), "Missing glibc loader"
    assert elf_count > 20, "Rootfs is unexpectedly empty"
    assert (root / "etc/s6-rc/source/boot/contents.d/dasungd").is_file(), "Dasung missing from base boot bundle"
    if configured:
        epoch = (root / "usr/share/fds/build-epoch").read_text().strip()
        assert epoch.isdecimal() and 0 < int(epoch) <= 4102444800, "Invalid image clock floor"
        assert not (root / "var/cache/ldconfig/aux-cache").exists(), "Build-specific linker auxiliary cache leaked into image"
        assert not list((root / "var/db/xbps").glob('*/aarch64-repodata')), "Build repository index leaked into image"
        for name in ("etc/ld.so.cache", "etc/udev/hwdb.bin", "etc/ssl/certs/ca-certificates.crt", "usr/lib/locale/locale-archive", "etc/s6-rc/compiled/db"):
            assert (root / name).stat().st_size > 0, f"Missing build-time cache: {name}"
        assert any((root / "var/cache/fontconfig").glob("*.cache-*")), "Missing build-time font cache"
        assert not any((root / "tmp").iterdir()), "Build files left in /tmp"
        assert not any((root / "run").iterdir()), "Runtime files leaked into image"
        init = rooted(root, "/sbin/init").read_text()
        assert init.startswith("#!/usr/bin/execlineb ") and "s6-linux-init " in init, "Init is not the native s6 launcher"
        assert '"/etc/s6-linux-init/current"' in init, "Wrong init configuration path"
        for service in ("runtime-fs", "hostname", "getty", "udev-trigger", "cartridged"):
            assert (root / f"etc/s6-rc/source/boot/contents.d/{service}").is_file(), f"Missing boot service: {service}"
        assert not (root / "etc/s6-rc/source/boot/contents.d/test-echo").exists(), "Test service must be opt-in"
        assert not (root / "usr/libexec/fds/m2-guest").exists(), "VM checker leaked into base rootfs"
    print(f"PASS: rootfs audit ({len(db)} packages, {elf_count} aarch64 ELF files, configured={configured})")


if __name__ == "__main__":
    try:
        audit(pathlib.Path(sys.argv[1]).resolve(), "--unconfigured" not in sys.argv[2:])
    except (AssertionError, OSError, ValueError, KeyError) as error:
        sys.exit(f"ERROR: {error}")