#!/usr/bin/env python3
"""Independently parse B17 images and compare frozen FreeBSD expectations."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
import shutil
import struct
import subprocess


SUPER_OFFSET = 1024
FEATURE_COMPAT_XATTR_FILTER = 0x00000004
FEATURE_COMPAT_PLAIN_XATTR_PFX = 0x00000010
FEATURE_INCOMPAT_XATTR_PREFIXES = 0x00000040
ACL_FILTER_BITS = (1 << 21) | (1 << 30)
CRC32C_POLY = 0x82F63B78


class Reject(Exception):
    def __init__(self, errno_name: str, point: str):
        super().__init__(f"{errno_name} at {point}")
        self.errno_name = errno_name
        self.point = point


def crc32c(data: bytes, seed: int = 0xFFFFFFFF) -> int:
    value = seed
    for byte in data:
        value ^= byte
        for _ in range(8):
            value = (value >> 1) ^ (CRC32C_POLY if value & 1 else 0)
    return value & 0xFFFFFFFF


class Reader:
    def __init__(self, path: Path):
        self.path = path
        self.data = path.read_bytes()
        if len(self.data) < SUPER_OFFSET + 144:
            raise Reject("EINTEGRITY", "super.bounds")
        if self.u32(SUPER_OFFSET) != 0xE0F5E1E2:
            raise Reject("EINTEGRITY", "super.magic")
        self.block_bits = self.data[SUPER_OFFSET + 12]
        if self.block_bits < 9 or self.block_bits > 16:
            raise Reject("EINTEGRITY", "super.block-size")
        self.block_size = 1 << self.block_bits
        self.blocks = self.u32(SUPER_OFFSET + 36)
        self.limit = self.blocks << self.block_bits
        if self.limit > len(self.data):
            raise Reject("EINTEGRITY", "super.image-bounds")
        self.feature_compat = self.u32(SUPER_OFFSET + 8)
        self.feature_incompat = self.u32(SUPER_OFFSET + 80)
        self.filter_reserved = self.data[SUPER_OFFSET + 104]
        self.meta_blkaddr = self.u32(SUPER_OFFSET + 40)
        self.xattr_blkaddr = self.u32(SUPER_OFFSET + 44)
        self.root_nid = self.u16(SUPER_OFFSET + 14)
        self.packed_nid = self.u64(SUPER_OFFSET + 96)
        self.prefix_count = self.data[SUPER_OFFSET + 91]
        self.prefix_start = self.u32(SUPER_OFFSET + 92)
        self.verify_checksum()

    def u16(self, offset: int) -> int:
        if offset < 0 or offset + 2 > len(self.data):
            raise Reject("EINTEGRITY", "raw.u16-bounds")
        return struct.unpack_from("<H", self.data, offset)[0]

    def u32(self, offset: int) -> int:
        if offset < 0 or offset + 4 > len(self.data):
            raise Reject("EINTEGRITY", "raw.u32-bounds")
        return struct.unpack_from("<I", self.data, offset)[0]

    def u64(self, offset: int) -> int:
        if offset < 0 or offset + 8 > len(self.data):
            raise Reject("EINTEGRITY", "raw.u64-bounds")
        return struct.unpack_from("<Q", self.data, offset)[0]

    def verify_checksum(self) -> None:
        if not self.feature_compat & 1:
            return
        end = self.block_size
        if end > len(self.data):
            raise Reject("EINTEGRITY", "super.checksum-bounds")
        expected = self.u32(SUPER_OFFSET + 4)
        block = bytearray(self.data[SUPER_OFFSET:end])
        block[4:8] = bytes(4)
        if crc32c(bytes(block)) != expected:
            raise Reject("EINTEGRITY", "super.checksum")

    def inode(self, nid: int) -> dict[str, int]:
        offset = (self.meta_blkaddr << self.block_bits) + (nid << 5)
        if offset + 32 > self.limit:
            raise Reject("EINTEGRITY", "inode.bounds")
        inode_format = self.u16(offset)
        inode_size = 64 if inode_format & 1 else 32
        if offset + inode_size > self.limit:
            raise Reject("EINTEGRITY", "inode.bounds")
        xattr_count = self.u16(offset + 2)
        xattr_size = 0 if xattr_count == 0 else 12 + 4 * (xattr_count - 1)
        size = self.u64(offset + 8) if inode_size == 64 else self.u32(offset + 8)
        return {
            "nid": nid,
            "offset": offset,
            "inode_size": inode_size,
            "xattr_count": xattr_count,
            "xattr_size": xattr_size,
            "layout": (inode_format >> 1) & 7,
            "size": size,
            "start_block": self.u32(offset + 16),
        }

    def inode_range(self, inode: dict[str, int], logical: int, length: int, point: str) -> bytes:
        if logical > inode["size"] or length > inode["size"] - logical:
            raise Reject("EINTEGRITY", point)
        if inode["layout"] == 2:
            physical = inode["offset"] + inode["inode_size"] + inode["xattr_size"] + logical
        elif inode["layout"] == 0:
            physical = (inode["start_block"] << self.block_bits) + logical
        else:
            raise Reject("EOPNOTSUPP", "inode.layout")
        if physical > self.limit or length > self.limit - physical:
            raise Reject("EINTEGRITY", point)
        return self.data[physical : physical + length]

    def directory_entries(self, inode: dict[str, int]) -> list[tuple[bytes, int]]:
        data = self.inode_range(inode, 0, inode["size"], "directory.bounds")
        if len(data) < 12:
            raise Reject("EINTEGRITY", "directory.header")
        first_name = struct.unpack_from("<H", data, 8)[0]
        if first_name == 0 or first_name % 12 != 0 or first_name > len(data):
            raise Reject("EINTEGRITY", "directory.name-offset")
        count = first_name // 12
        entries = []
        for index in range(count):
            entry = index * 12
            nid = struct.unpack_from("<Q", data, entry)[0]
            name_start = struct.unpack_from("<H", data, entry + 8)[0]
            name_end = (
                struct.unpack_from("<H", data, entry + 20)[0]
                if index + 1 < count
                else len(data)
            )
            if name_start > name_end or name_end > len(data):
                raise Reject("EINTEGRITY", "directory.name-bounds")
            name = data[name_start:name_end].split(b"\0", 1)[0]
            entries.append((name, nid))
        return entries

    def resolve(self, path: str) -> dict[str, int]:
        inode = self.inode(self.root_nid)
        for component in path.strip("/").encode("ascii").split(b"/"):
            if not component:
                continue
            nid = next(
                (candidate for name, candidate in self.directory_entries(inode)
                 if name == component),
                None,
            )
            if nid is None:
                raise Reject("ENOATTR", "path.lookup")
            inode = self.inode(nid)
        return inode

    def load_prefixes(self) -> list[tuple[int, bytes]]:
        if not self.feature_incompat & FEATURE_INCOMPAT_XATTR_PREFIXES or self.prefix_count == 0:
            return []
        if self.feature_compat & FEATURE_COMPAT_PLAIN_XATTR_PFX:
            raise Reject("EOPNOTSUPP", "prefix.plain-unexpected")
        if self.packed_nid == 0:
            raise Reject("EINTEGRITY", "prefix.carrier")
        packed = self.inode(self.packed_nid)
        logical = self.prefix_start << 2
        prefixes = []
        for _ in range(self.prefix_count):
            while logical % 4:
                logical += 1
            header = self.inode_range(packed, logical, 2, "prefix.header-bounds")
            length = struct.unpack("<H", header)[0]
            if length < 1 or length > 256:
                raise Reject("EINTEGRITY", "prefix.length")
            payload = self.inode_range(
                packed, logical + 2, length, "prefix.payload-bounds"
            )
            base_index = payload[0]
            infix = payload[1:]
            if b"\0" in infix:
                raise Reject("EINTEGRITY", "prefix.infix-nul")
            prefixes.append((base_index, infix))
            logical += 2 + length
        return prefixes

    def load_body(self, inode: dict[str, int]) -> tuple[int, int, list[int]]:
        size = inode["xattr_size"]
        if size == 0:
            return 0, 0, []
        body = inode["offset"] + inode["inode_size"]
        if body > self.limit or size > self.limit - body:
            raise Reject("EINTEGRITY", "ibody.bounds")
        if size < 12:
            raise Reject("EINTEGRITY", "ibody.header-bounds")
        if size == 12:
            raise Reject("EOPNOTSUPP", "ibody.header-only")
        shared_count = self.data[body + 4]
        header_size = 12 + shared_count * 4
        if header_size > size:
            raise Reject("EINTEGRITY", "ibody.shared-count")
        shared = [self.u32(body + 12 + index * 4) for index in range(shared_count)]
        return body, header_size, shared

    def entry(self, offset: int, limit: int, kind: str) -> dict[str, object]:
        if offset > limit or 4 > limit - offset:
            raise Reject("EINTEGRITY", f"{kind}.entry-header")
        name_length = self.data[offset]
        name_index = self.data[offset + 1]
        value_length = self.u16(offset + 2)
        name_end = offset + 4 + name_length
        if name_end > limit:
            raise Reject("EINTEGRITY", f"{kind}.name-bounds")
        value_end = name_end + value_length
        if value_end > limit:
            raise Reject("EINTEGRITY", f"{kind}.value-bounds")
        aligned_end = (value_end + 3) & ~3
        if aligned_end > limit:
            raise Reject("EINTEGRITY", f"{kind}.padding-bounds")
        name = self.data[offset + 4 : name_end]
        if b"\0" in name:
            raise Reject("EINTEGRITY", f"{kind}.name-nul")
        return {
            "offset": offset,
            "next": aligned_end,
            "name_index": name_index,
            "name": name,
            "value": self.data[name_end:value_end],
        }

    def entries(self, inode: dict[str, int]) -> tuple[list[dict[str, object]], int]:
        body, header_size, shared_ids = self.load_body(inode)
        if body == 0:
            return [], 0
        body_end = body + inode["xattr_size"]
        cursor = body + header_size
        entries = []
        while cursor < body_end:
            item = self.entry(cursor, body_end, "inline")
            entries.append(item)
            cursor = int(item["next"])
        for shared_id in shared_ids:
            offset = (self.xattr_blkaddr << self.block_bits) + shared_id * 4
            if offset > self.limit or 4 > self.limit - offset:
                raise Reject("EINTEGRITY", "shared.offset-bounds")
            entries.append(self.entry(offset, self.limit, "shared"))
        return entries, len(shared_ids)

    def resolved_name(
        self, entry: dict[str, object], prefixes: list[tuple[int, bytes]]
    ) -> tuple[str, bytes]:
        index = int(entry["name_index"])
        infix = b""
        from_prefix = False
        if index & 0x80:
            prefix_id = index & 0x7F
            if prefix_id >= len(prefixes):
                raise Reject("ENOATTR", "name.long-prefix-id")
            index, infix = prefixes[prefix_id]
            from_prefix = True
        mapping = {
            1: ("user", b""),
            2: ("system", b"posix_acl_access"),
            3: ("system", b"posix_acl_default"),
            4: ("system", b"trusted."),
            6: ("system", b"security."),
        }
        if index not in mapping:
            point = "name.prefix-base-index" if from_prefix else "name.short-index"
            raise Reject("ENOATTR", point)
        namespace, fixed = mapping[index]
        return namespace, fixed + infix + bytes(entry["name"])

    def getxattr(self, inode: dict[str, int], namespace: str, name: bytes,
                 prefixes: list[tuple[int, bytes]]) -> tuple[bytes, int]:
        entries, shared_count = self.entries(inode)
        deferred: Reject | None = None
        for entry in entries:
            try:
                actual_namespace, actual_name = self.resolved_name(entry, prefixes)
            except Reject as error:
                deferred = error
                continue
            if actual_namespace == namespace and actual_name == name:
                return bytes(entry["value"]), shared_count
        if deferred is not None:
            raise deferred
        raise Reject("ENOATTR", "name.not-found")

    def listxattr(self, inode: dict[str, int], namespace: str,
                  prefixes: list[tuple[int, bytes]]) -> tuple[list[bytes], int]:
        entries, shared_count = self.entries(inode)
        names = []
        deferred: Reject | None = None
        for entry in entries:
            try:
                actual_namespace, actual_name = self.resolved_name(entry, prefixes)
            except Reject as error:
                deferred = error
                continue
            if actual_namespace == namespace:
                names.append(actual_name)
        if deferred is not None and not names:
            raise deferred
        return names, shared_count


def parse_acl(value: bytes) -> list[list[int]]:
    if len(value) < 4 or (len(value) - 4) % 8:
        raise Reject("EINTEGRITY", "acl.value-size")
    if struct.unpack_from("<I", value)[0] != 2:
        raise Reject("EINTEGRITY", "acl.version")
    entries = []
    for offset in range(4, len(value), 8):
        entries.append(list(struct.unpack_from("<HHI", value, offset)))
    return entries


def execute(case: dict, image_path: Path) -> dict[str, object]:
    reader = Reader(image_path)
    prefixes = reader.load_prefixes()
    if "expected_feature_filter" in case:
        actual_feature = bool(reader.feature_compat & FEATURE_COMPAT_XATTR_FILTER)
        if actual_feature != case["expected_feature_filter"]:
            raise AssertionError("raw xattr filter feature differs")
        if reader.filter_reserved != case["expected_filter_reserved"]:
            raise AssertionError("raw xattr filter reserved byte differs")
        usable = actual_feature and reader.filter_reserved == 0
        if usable != case["expected_filter_usable"]:
            raise AssertionError("xattr filter use-site gate differs")
    operation = case["operation"]
    result: dict[str, object] = {}
    if operation == "mount":
        return result
    inode = reader.resolve(case["target"])
    if "expected_acl_filter_negative" in case:
        body, _, _ = reader.load_body(inode)
        actual_negative = (
            reader.u32(body) & ACL_FILTER_BITS
        ) == ACL_FILTER_BITS
        if actual_negative != case["expected_acl_filter_negative"]:
            raise AssertionError("ACL name-filter declaration differs")
    if operation == "get":
        value, shared_count = reader.getxattr(
            inode, case["namespace"], case["name"].encode("ascii"), prefixes
        )
        result["value_hex"] = value.hex()
        result["shared_count"] = shared_count
        if "expected_value_hex" in case and value.hex() != case["expected_value_hex"]:
            raise AssertionError("xattr value differs")
        if "expected_shared_count" in case and shared_count != case["expected_shared_count"]:
            raise AssertionError("xattr shared count differs")
    elif operation == "list":
        names, shared_count = reader.listxattr(inode, case["namespace"], prefixes)
        result["names"] = [name.decode("ascii") for name in names]
        result["shared_count"] = shared_count
    elif operation == "acl":
        value, _ = reader.getxattr(
            inode, "system", b"posix_acl_access", prefixes
        )
        acl = parse_acl(value)
        result["acl"] = acl
        if acl != case["expected_acl"]:
            raise AssertionError("ACL empty-suffix value differs")
    else:
        raise AssertionError(f"unknown operation: {operation}")
    return result


def fsck_accept(path: Path) -> dict[str, object]:
    fsck = shutil.which("fsck.erofs")
    if fsck is None:
        raise SystemExit("fsck.erofs is required for legal B17 fixture baseline")
    completed = subprocess.run(
        [fsck, "-d0", "--xattrs", str(path)], check=False, text=True,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    )
    return {"exit": completed.returncode, "stdout": completed.stdout}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--spec", type=Path, required=True)
    parser.add_argument("--fixtures", type=Path, required=True)
    parser.add_argument("--report", type=Path, required=True)
    args = parser.parse_args()
    spec = json.loads(args.spec.read_text(encoding="ascii"))
    index = json.loads(
        (args.fixtures / "fixture-index.json").read_text(encoding="ascii")
    )
    indexed = {item["id"]: item for item in index["fixtures"]}
    errnos = spec["freebsd_errno"]
    results = []
    failures = []
    for case in spec["cases"]:
        image_path = args.fixtures / indexed[case["id"]]["path"]
        actual_errno = 0
        actual_errno_name = None
        actual_reject = "accepted"
        details: dict[str, object] = {}
        try:
            details = execute(case, image_path)
        except Reject as error:
            actual_errno_name = error.errno_name
            actual_errno = errnos[error.errno_name]
            actual_reject = error.point
        except Exception as error:
            failures.append({"id": case["id"], "reason": repr(error)})
            results.append({"id": case["id"], "status": "FAIL", "exception": repr(error)})
            continue
        expected_name = case.get("expected_errno_name")
        matched = (
            actual_errno == case["expected_errno"]
            and actual_errno_name == expected_name
            and actual_reject == case["expected_reject"]
        )
        fsck = None
        if case.get("fsck_accept"):
            fsck = fsck_accept(image_path)
            matched = matched and fsck["exit"] == 0
        result = {
            "id": case["id"],
            "class": case["class"],
            "status": "PASS" if matched else "FAIL",
            "expected_errno": case["expected_errno"],
            "actual_errno": actual_errno,
            "actual_errno_name": actual_errno_name,
            "expected_reject": case["expected_reject"],
            "actual_reject": actual_reject,
            "details": details,
        }
        if fsck is not None:
            result["fsck_exit"] = fsck["exit"]
            result["fsck_stdout_sha256"] = hashlib.sha256(
                fsck["stdout"].encode("utf-8")
            ).hexdigest()
        results.append(result)
        if not matched:
            failures.append(result)
    report = {
        "schema": 1,
        "batch": "B17",
        "status": "PASS" if not failures else "FAIL",
        "fixture_count": len(results),
        "legal_passed": sum(
            item["status"] == "PASS" and item["class"] == "legal"
            for item in results
        ),
        "damaged_passed": sum(
            item["status"] == "PASS" and item["class"] == "damaged"
            for item in results
        ),
        "failures": failures,
        "results": results,
    }
    args.report.write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="ascii"
    )
    print(
        json.dumps(
            {
                "status": report["status"],
                "fixtures": report["fixture_count"],
                "legal_passed": report["legal_passed"],
                "damaged_passed": report["damaged_passed"],
            },
            sort_keys=True,
        )
    )
    return 0 if not failures else 1


if __name__ == "__main__":
    raise SystemExit(main())
