#!/usr/bin/env python3
"""Validate release-helper format, architecture, and Linux linkage contract."""

from __future__ import annotations

from pathlib import Path
import os
import struct
import sys
from typing import BinaryIO


PE_TARGET = "x86_64-pc-windows-msvc"
GNU_TARGET = "x86_64-unknown-linux-gnu"
MUSL_TARGET = "x86_64-unknown-linux-musl"


class ValidationError(ValueError):
    """The candidate cannot satisfy its labeled release target."""


def read_exact_at(candidate: BinaryIO, offset: int, size: int, label: str) -> bytes:
    candidate.seek(offset)
    value = candidate.read(size)
    if len(value) != size:
        raise ValidationError(f"truncated {label}")
    return value


def validate_pe_x86_64(candidate: BinaryIO, size: int) -> str:
    dos = read_exact_at(candidate, 0, 64, "DOS header")
    if dos[:2] != b"MZ":
        raise ValidationError("not a PE executable (missing MZ magic)")

    pe_offset = struct.unpack_from("<I", dos, 0x3C)[0]
    coff_size = 24
    if pe_offset < len(dos) or pe_offset > size - coff_size:
        raise ValidationError("invalid PE header offset")
    coff = read_exact_at(candidate, pe_offset, coff_size, "PE/COFF header")
    if coff[:4] != b"PE\0\0":
        raise ValidationError("not a PE executable (missing PE signature)")

    machine, section_count = struct.unpack_from("<HH", coff, 4)
    optional_size, characteristics = struct.unpack_from("<HH", coff, 20)
    if machine != 0x8664:
        raise ValidationError(
            f"PE machine 0x{machine:04x} is not x86_64 (expected 0x8664)"
        )
    if section_count == 0:
        raise ValidationError("PE image has no sections")
    if characteristics & 0x0002 == 0:
        raise ValidationError("PE image is not marked executable")
    if characteristics & 0x2000:
        raise ValidationError("PE image is marked as a DLL")

    optional_offset = pe_offset + coff_size
    fixed_optional_size = 112
    if (
        optional_size < fixed_optional_size
        or optional_offset > size
        or optional_size > size - optional_offset
    ):
        raise ValidationError("invalid PE32+ optional header size")
    optional = read_exact_at(
        candidate, optional_offset, optional_size, "PE32+ optional header"
    )
    optional_magic = struct.unpack_from("<H", optional)[0]
    if optional_magic != 0x20B:
        raise ValidationError(
            f"PE optional-header magic 0x{optional_magic:04x} is not PE32+"
        )
    entry_point = struct.unpack_from("<I", optional, 16)[0]
    directory_count = struct.unpack_from("<I", optional, 108)[0]
    if directory_count > (optional_size - fixed_optional_size) // 8:
        raise ValidationError(
            "PE32+ optional header is truncated for its data directories"
        )
    if entry_point == 0:
        raise ValidationError("PE image has zero entry point")

    section_table_offset = optional_offset + optional_size
    section_table_size = section_count * 40
    if (
        section_table_offset > size
        or section_table_size > size - section_table_offset
    ):
        raise ValidationError("PE section table is outside the file")
    sections = read_exact_at(
        candidate, section_table_offset, section_table_size, "PE section table"
    )

    entry_mapped = False
    entry_executable = False
    for index in range(section_count):
        section_offset = index * 40
        (
            virtual_size,
            virtual_address,
            raw_size,
            raw_offset,
        ) = struct.unpack_from("<IIII", sections, section_offset + 8)
        section_characteristics = struct.unpack_from(
            "<I", sections, section_offset + 36
        )[0]
        if raw_offset > size or raw_size > size - raw_offset:
            raise ValidationError("PE section raw data is outside the file")

        mapped_size = max(virtual_size, raw_size)
        if mapped_size > 0x1_0000_0000 - virtual_address:
            raise ValidationError("PE section virtual range overflows the RVA space")
        if virtual_address <= entry_point < virtual_address + mapped_size:
            entry_mapped = True
            if section_characteristics & 0x20000000:
                entry_executable = True

    if not entry_mapped:
        raise ValidationError("PE entry point is not mapped by any section")
    if not entry_executable:
        raise ValidationError("PE entry point is in a non-executable section")
    return "PE32+ x86_64 executable"


def validate_elf_x86_64(candidate: BinaryIO, size: int, target: str) -> str:
    header = read_exact_at(candidate, 0, 64, "ELF64 header")
    if header[:4] != b"\x7fELF":
        raise ValidationError("not an ELF executable (missing ELF magic)")
    if header[4] != 2:
        raise ValidationError(f"ELF class {header[4]} is not 64-bit")
    if header[5] != 1:
        raise ValidationError(f"ELF data encoding {header[5]} is not little-endian")
    if header[6] != 1:
        raise ValidationError(f"ELF identification version {header[6]} is invalid")

    (
        elf_type,
        machine,
        version,
        entry_point,
        program_offset,
        _section_offset,
        _flags,
        header_size,
        program_entry_size,
        program_count,
        _section_entry_size,
        _section_count,
        _section_names,
    ) = struct.unpack_from("<HHIQQQIHHHHHH", header, 16)
    if elf_type not in (2, 3):
        raise ValidationError(f"ELF type {elf_type} is not executable or PIE")
    if machine != 62:
        raise ValidationError(f"ELF machine {machine} is not x86_64 (expected 62)")
    if version != 1 or header_size != 64:
        raise ValidationError("invalid ELF64 header version or size")
    if program_entry_size != 56 or program_count == 0:
        raise ValidationError("ELF64 program-header table is missing or malformed")
    table_size = program_entry_size * program_count
    if program_offset < header_size or program_offset > size:
        raise ValidationError("ELF64 program-header table is outside the file")
    if table_size > size - program_offset:
        raise ValidationError("ELF64 program-header table is outside the file")

    interpreters: list[str] = []
    needed_count = 0
    has_load = False
    entry_mapped = False
    entry_executable = False
    for index in range(program_count):
        offset = program_offset + index * program_entry_size
        program = read_exact_at(candidate, offset, program_entry_size, "ELF64 program header")
        (
            program_type,
            program_flags,
            file_offset,
            virtual_address,
            _physical_address,
            file_size,
            memory_size,
            _alignment,
        ) = struct.unpack("<IIQQQQQQ", program)
        if file_offset > size or file_size > size - file_offset:
            raise ValidationError("ELF64 segment is outside the file")
        if program_type == 1:
            has_load = True
            if file_size > memory_size:
                raise ValidationError(
                    "ELF64 load segment file size exceeds memory size"
                )
            if memory_size > 0x1_0000_0000_0000_0000 - virtual_address:
                raise ValidationError("ELF64 load segment virtual range overflows")
            if virtual_address <= entry_point < virtual_address + memory_size:
                entry_mapped = True
                if program_flags & 0x1:
                    entry_executable = True
        elif program_type == 3:
            if file_size < 2 or file_size > 4096:
                raise ValidationError("ELF interpreter segment has an invalid size")
            raw_interpreter = read_exact_at(
                candidate, file_offset, file_size, "ELF interpreter"
            )
            if raw_interpreter[-1:] != b"\0":
                raise ValidationError("ELF interpreter is not NUL-terminated")
            try:
                interpreters.append(raw_interpreter[:-1].decode("ascii"))
            except UnicodeDecodeError as error:
                raise ValidationError("ELF interpreter is not ASCII") from error
        elif program_type == 2:
            if file_size % 16 != 0:
                raise ValidationError("ELF dynamic segment has a malformed size")
            for dynamic_offset in range(file_offset, file_offset + file_size, 16):
                tag, _value = struct.unpack(
                    "<qQ", read_exact_at(candidate, dynamic_offset, 16, "ELF dynamic entry")
                )
                if tag == 0:
                    break
                if tag == 1:
                    needed_count += 1

    if not has_load:
        raise ValidationError("ELF image has no loadable segment")
    if entry_point == 0:
        raise ValidationError("ELF image has zero entry point")
    if not entry_mapped:
        raise ValidationError("ELF entry point is not mapped by a loadable segment")
    if not entry_executable:
        raise ValidationError(
            "ELF entry point is in a non-executable loadable segment"
        )
    if target == GNU_TARGET:
        if len(interpreters) != 1:
            raise ValidationError(
                "GNU target must declare exactly one glibc x86_64 program interpreter"
            )
        interpreter = interpreters[0]
        if os.path.basename(interpreter) != "ld-linux-x86-64.so.2":
            raise ValidationError(
                f"GNU target has unexpected program interpreter {interpreter!r}"
            )
        return f"ELF64 x86_64 GNU-linked executable ({interpreter})"

    if interpreters:
        raise ValidationError(
            f"musl target must be static but declares interpreter {interpreters[0]!r}"
        )
    if needed_count:
        raise ValidationError(
            f"musl target must be static but has {needed_count} DT_NEEDED entr{'y' if needed_count == 1 else 'ies'}"
        )
    return "ELF64 x86_64 static executable (no interpreter or DT_NEEDED entries)"


def validate(path: Path, target: str) -> str:
    if not path.is_file():
        raise ValidationError("candidate is not a regular file")
    size = path.stat().st_size
    with path.open("rb") as candidate:
        if target == PE_TARGET:
            return validate_pe_x86_64(candidate, size)
        if target in (GNU_TARGET, MUSL_TARGET):
            return validate_elf_x86_64(candidate, size, target)
    raise ValidationError(f"unsupported release target {target!r}")


def main() -> int:
    if len(sys.argv) != 3:
        print(
            "usage: validate-release-binary.py <target-triple> <binary>",
            file=sys.stderr,
        )
        return 2
    target = sys.argv[1]
    path = Path(sys.argv[2])
    try:
        description = validate(path, target)
    except (OSError, ValidationError) as error:
        print(f"INVALID {target}: {path}: {error}", file=sys.stderr)
        return 1
    print(f"ok   {target}: {path} ({description})")
    return 0


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