#!/overlayer/syshub/bin/python3
# Follow-up release tests: module signatures under lockdown, and fs-verity digests.
# Run as root: elevate python3 -I ~/zainium-test/release-tests-2.py
import ctypes, errno, os, struct, subprocess, time

LOG = "/home/alizain/zainium-test/release-tests-2.log"
out = open(LOG, "w")
results = []

def log(s=""):
    print(s); out.write(s + "\n"); out.flush()

def check(name, ok, detail=""):
    results.append((name, ok))
    log(f"[{'PASS' if ok else 'FAIL'}] {name}" + (f" -- {detail}" if detail else ""))

libc = ctypes.CDLL(None, use_errno=True)
SYS_finit_module, SYS_delete_module = 313, 176
MOD = "/overlayer/syshub/drivers/modules/linux-zainium-7.3/kernel/drivers/net/dummy.ko"
MAGIC = b"~Module signature appended~\n"

def loaded(name):
    return any(l.split(" ")[0] == name for l in open("/proc/modules"))

def finit(path):
    fd = os.open(path, os.O_RDONLY)
    r = libc.syscall(SYS_finit_module, fd, b"", 0)
    e = ctypes.get_errno()
    os.close(fd)
    return r, e

def unload(name):
    libc.syscall(SYS_delete_module, name.encode(), os.O_NONBLOCK)

log(f"Zainium OS release tests 2 -- {time.strftime('%Y-%m-%d %H:%M:%S %z')}")
log(f"kernel: {os.uname().release} {os.uname().version}")
log(f"lockdown: {open('/sys/kernel/security/lockdown').read().strip()}")

log("\n== Module signatures ==")
if loaded("dummy"):
    log("dummy is already loaded; skipping the module tests")
else:
    data = open(MOD, "rb").read()
    check("the shipped dummy.ko carries a signature", data.endswith(MAGIC))
    # struct module_signature (12 bytes) sits right before the magic; its last 4 bytes are the signature length
    sig_len = struct.unpack(">I", data[-len(MAGIC) - 4:-len(MAGIC)])[0]
    stripped = data[:len(data) - len(MAGIC) - 12 - sig_len]
    tmp = "/run/zainium-unsigned-dummy.ko"
    open(tmp, "wb").write(stripped)
    r, e = finit(tmp)
    os.unlink(tmp)
    if r == 0:
        unload("dummy")
    check("an unsigned module is refused (lockdown=integrity)", r != 0,
          "LOADED" if r == 0 else f"refused: {errno.errorcode.get(e, e)}")
    tampered = bytearray(data)
    tampered[len(stripped) // 2] ^= 0xFF  # one changed byte inside the signed part
    tmp = "/run/zainium-tampered-dummy.ko"
    open(tmp, "wb").write(bytes(tampered))
    r, e = finit(tmp)
    os.unlink(tmp)
    if r == 0:
        unload("dummy")
    check("a module changed after signing is refused", r != 0,
          "LOADED" if r == 0 else f"refused: {errno.errorcode.get(e, e)}")
    r, e = finit(MOD)
    ok = r == 0 or e == errno.EEXIST
    check("the original signed module loads", ok, "loaded" if r == 0 else errno.errorcode.get(e, e))
    if r == 0:
        unload("dummy")

log("\n== fs-verity digests (fs-verity digest, SHA-256, 4096-byte blocks) ==")
for p in ["/overlayer/syshub/bin/elevate", "/overlayer/syshub/bin/zex", "/overlayer/syshub/bin/sudo",
          "/overlayer/syshub/lib/libsys_resolver.so", "/overlayer/syshub/lib/systemd/systemd",
          "/overlayer/zaisys/kernel/linux-zainium", "/overlayer/zaisys/kernel/quantra-ramfs.img"]:
    r = subprocess.run(["/home/alizain/.cargo/bin/fs-verity", "digest", p], capture_output=True, text=True)
    log("  " + (r.stdout.strip() or r.stderr.strip()))

log("\n== Summary ==")
passed = sum(1 for _, ok in results if ok)
log(f"{passed}/{len(results)} checks passed")
for n, ok in results:
    if not ok:
        log(f"  FAILED: {n}")
out.close()
os.chmod(LOG, 0o644)