#!/usr/bin/env python3
"""Build an ISO-only FIPS package without its conflicting default nft conffile.

This does not install a package or change an ISO. The input SHA is mandatory so a
different upstream package cannot silently inherit this narrowly reviewed edit.
"""

import argparse
import copy
import hashlib
import io
import json
import os
from pathlib import Path
import tarfile

REVIEWED_INPUTS = {
    "c2a3a77babdae63297104c32e3028fef89c80080f753893749c2c9018497d2e6":
        ("0.4.1-1", "0.4.1-1+archy.iso1"),
    "69a029196555a5d729a4c9e4448e319bc5b927cbb33859acff20fa651ea047d7":
        ("0.5.3+archy.handoff1", "0.5.3+archy.handoff1.iso1"),
}
REMOVED_DATA = "./etc/fips/fips.nft"
REMOVED_CONFFILE = "/etc/fips/fips.nft"
AR_MAGIC = b"!<arch>\n"


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def read_ar(data):
    if not data.startswith(AR_MAGIC):
        raise ValueError("not an ar archive")
    members = []
    pos = len(AR_MAGIC)
    while pos < len(data):
        header = data[pos:pos + 60]
        if len(header) != 60 or header[58:] != b"`\n":
            raise ValueError("invalid ar member header")
        name = header[:16].decode("ascii").strip().rstrip("/")
        size = int(header[48:58].decode("ascii").strip())
        pos += 60
        payload = data[pos:pos + size]
        if len(payload) != size:
            raise ValueError("truncated ar member")
        members.append((name, payload))
        pos += size + (size % 2)
    if pos != len(data):
        raise ValueError("invalid ar padding")
    return members


def write_ar(members):
    output = bytearray(AR_MAGIC)
    for name, data in members:
        if len(name) > 15:
            raise ValueError("ar member name too long")
        header = f"{name + '/':<16}{0:<12}{0:<6}{0:<6}{'100644':<8}{len(data):<10}`\n".encode("ascii")
        if len(header) != 60:
            raise AssertionError("ar header width")
        output.extend(header)
        output.extend(data)
        if len(data) % 2:
            output.extend(b"\n")
    return bytes(output)


def tar_items(blob):
    with tarfile.open(fileobj=io.BytesIO(blob), mode="r:xz") as archive:
        return [(member, archive.extractfile(member).read() if member.isfile() else None)
                for member in archive.getmembers()]


def tar_bytes(items):
    buffer = io.BytesIO()
    with tarfile.open(fileobj=buffer, mode="w:xz", format=tarfile.GNU_FORMAT) as archive:
        for member, data in items:
            archive.addfile(member, io.BytesIO(data) if data is not None else None)
    return buffer.getvalue()


def normalized(name):
    return name.lstrip("./")


def derive(original, old_version="0.4.1-1", new_version="0.4.1-1+archy.iso1"):
    members = read_ar(original)
    if [name for name, _ in members] != ["debian-binary", "control.tar.xz", "data.tar.xz"]:
        raise ValueError("unexpected deb member layout")
    if members[0][1] != b"2.0\n":
        raise ValueError("unexpected deb format")
    controls = tar_items(members[1][1])
    data = tar_items(members[2][1])
    names = [normalized(item.name) for item, _ in data]
    if names.count(normalized(REMOVED_DATA)) != 1:
        raise ValueError("expected exactly one default fips.nft")
    if len(names) != len(set(names)):
        raise ValueError("duplicate data path")
    control_names = [normalized(item.name) for item, _ in controls]
    if control_names.count("control") != 1 or control_names.count("conffiles") != 1:
        raise ValueError("unexpected control members")
    if "md5sums" in control_names:
        # Neither reviewed input contains this member. A future input with one
        # needs a separate, explicit digest-entry transformation and review.
        raise ValueError("md5sums requires a separately reviewed transformation")
    changed = {}
    updated_controls = []
    for item, content in controls:
        name = normalized(item.name)
        if name == "control":
            if content.count(b"Package: fips\n") != 1 or content.count(b"Architecture: amd64\n") != 1:
                raise ValueError("unexpected package or architecture")
            old = f"Version: {old_version}\n".encode()
            if content.count(old) != 1:
                raise ValueError("unexpected package version")
            replacement = content.replace(old, f"Version: {new_version}\n".encode())
            changed[name] = {"before": sha256(content), "after": sha256(replacement)}
            content = replacement
        elif name == "conffiles":
            line = (REMOVED_CONFFILE + "\n").encode()
            if content.count(line) != 1:
                raise ValueError("unexpected conffile declaration")
            replacement = content.replace(line, b"")
            changed[name] = {"before": sha256(content), "after": sha256(replacement)}
            content = replacement
        updated_item = copy.copy(item)
        if name in changed:
            updated_item.size = len(content)
        updated_controls.append((updated_item, content))
    updated_data = [(item, content) for item, content in data
                    if normalized(item.name) != normalized(REMOVED_DATA)]
    if len(updated_data) != len(data) - 1:
        raise AssertionError("data removal count")
    output = write_ar([
        members[0],
        (members[1][0], tar_bytes(updated_controls)),
        (members[2][0], tar_bytes(updated_data)),
    ])
    rebuilt = read_ar(output)
    new_controls, new_data = tar_items(rebuilt[1][1]), tar_items(rebuilt[2][1])
    for before, after in ((updated_controls, new_controls), (updated_data, new_data)):
        if len(before) != len(after):
            raise AssertionError("member count changed")
        for (old_member, old_content), (new_member, new_content) in zip(before, after):
            old_info, new_info = old_member.get_info(), new_member.get_info()
            # tar header checksums vary with serialization; all semantic metadata
            # and unchanged file bytes must still match exactly.
            old_info.pop("chksum")
            new_info.pop("chksum")
            if (old_info != new_info or old_member.pax_headers != new_member.pax_headers
                    or old_content != new_content):
                raise AssertionError(f"member metadata or content changed: {old_member.name}: {old_info} != {new_info}; data={old_content == new_content}")
    return output, {
        "schema": "iso-fips-derivative-v1",
        "original_sha256": sha256(original),
        "derivative_sha256": sha256(output),
        "old_version": old_version,
        "new_version": new_version,
        "removed_data_path": normalized(REMOVED_DATA),
        "removed_data_sha256": sha256(next(content for item, content in data if normalized(item.name) == normalized(REMOVED_DATA))),
        "changed_control_members": changed,
        "unchanged_data_members": len(updated_data),
        "unchanged_control_members": len(controls) - 2,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--expected-input-sha256", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--receipt", type=Path, required=True)
    args = parser.parse_args()
    original = args.input.read_bytes()
    if sha256(original) != args.expected_input_sha256:
        raise SystemExit("input SHA mismatch")
    if args.expected_input_sha256 not in REVIEWED_INPUTS:
        raise SystemExit("input SHA is not a reviewed FIPS package")
    output, receipt = derive(original, *REVIEWED_INPUTS[args.expected_input_sha256])
    args.output.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    args.receipt.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    if any(directory.stat().st_mode & 0o077 for directory in
           {args.output.parent, args.receipt.parent}):
        raise SystemExit("private output directories must be mode 0700")
    receipt_bytes = (json.dumps(receipt, indent=2, sort_keys=True) + "\n").encode()
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    output_fd = os.open(args.output, flags, 0o600)
    try:
        receipt_fd = os.open(args.receipt, flags, 0o600)
    except BaseException:
        os.close(output_fd)
        raise
    with os.fdopen(output_fd, "wb") as handle:
        handle.write(output)
        handle.flush()
        os.fsync(handle.fileno())
    with os.fdopen(receipt_fd, "wb") as handle:
        handle.write(receipt_bytes)
        handle.flush()
        os.fsync(handle.fileno())
    for directory in {args.output.parent, args.receipt.parent}:
        fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)


if __name__ == "__main__":
    main()