#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Synthetic byte tests for the Phase-1.0AG ELF admission contract."""

from __future__ import annotations

import argparse
import hashlib
from pathlib import Path
import struct
import sys
import unittest


PARSER = argparse.ArgumentParser()
PARSER.add_argument("--root", type=Path, required=True)
ROOT = PARSER.parse_args().root
sys.path.insert(0, str(ROOT / "tools"))

from phase10ag_bounded_elf import *  # noqa: E402,F403


def make_elf(headers: list[tuple[int, int, int, int, int, int, int, int]] | None = None,
             entry: int = 0x1000, elf_type: int = ET_DYN,
             machine: int = EM_X86_64) -> bytes:
    if headers is None:
        headers = [(PT_LOAD, PF_R | PF_X, 0x1000, 0x1000, 0, 0x100, 0x100, 0x1000),
                   (PT_LOAD, PF_R | PF_W, 0x2000, 0x3000, 0, 0x80, 0x180, 0x1000)]
    size = max([0x3000] + [item[2] + item[5] for item in headers])
    payload = bytearray(size)
    ident = bytearray(16)
    ident[:7] = b"\x7fELF\x02\x01\x01"
    ELF_HEADER.pack_into(payload, 0, bytes(ident), elf_type, machine, 1, entry,
                         64, 0, 0, 64, 56, len(headers), 64, 0, 0)
    for index, header in enumerate(headers):
        PROGRAM_HEADER.pack_into(payload, 64 + index * 56, *header)
    return bytes(payload)


def assess(payload: bytes):
    return assess_elf(payload, hashlib.sha256(payload).hexdigest())


class BoundedElfTests(unittest.TestCase):
    def test_valid_pie(self) -> None:
        result = assess(make_elf())
        self.assertEqual(len(result.load_segments), 2)
        self.assertFalse(result.writable_executable_segment)
        self.assertFalse(result.execution_performed)

    def test_hash_must_match(self) -> None:
        with self.assertRaises(ElfContractError):
            assess_elf(make_elf(), "0" * 64)

    def test_header_identity_type_and_machine(self) -> None:
        for payload in (b"not-elf" + b"\0" * 100, make_elf(elf_type=2),
                        make_elf(machine=3)):
            with self.subTest(), self.assertRaises(ElfContractError):
                assess(payload)

    def test_rejects_interpreter(self) -> None:
        headers = [(PT_INTERP, PF_R, 0x300, 0x300, 0, 8, 8, 1),
                   (PT_LOAD, PF_R | PF_X, 0x1000, 0x1000, 0, 0x100, 0x100, 0x1000)]
        with self.assertRaises(ElfContractError):
            assess(make_elf(headers))

    def test_rejects_writable_executable(self) -> None:
        headers = [(PT_LOAD, PF_R | PF_W | PF_X, 0x1000, 0x1000, 0,
                    0x100, 0x100, 0x1000)]
        with self.assertRaises(ElfContractError):
            assess(make_elf(headers))

    def test_rejects_file_and_memory_bounds(self) -> None:
        cases = [
            [(PT_LOAD, PF_R | PF_X, 0x2f80, 0x1000, 0, 0x100, 0x100, 0x1000)],
            [(PT_LOAD, PF_R | PF_X, 0x1000, 0x1000, 0, 0x200, 0x100, 0x1000)],
            [(PT_LOAD, PF_R | PF_X, 0x1000, 0x1000, 0, 0x100,
              MAX_TOTAL_LOAD_MEMORY + 1, 0x1000)],
        ]
        for headers in cases:
            with self.subTest(headers=headers), self.assertRaises(ElfContractError):
                assess(make_elf(headers))

    def test_rejects_bad_alignment(self) -> None:
        for alignment in (0, 3, MAX_ALIGNMENT * 2):
            headers = [(PT_LOAD, PF_R | PF_X, 0x1000, 0x1000, 0,
                        0x100, 0x100, alignment)]
            with self.subTest(alignment=alignment), self.assertRaises(ElfContractError):
                assess(make_elf(headers))

    def test_rejects_overlapping_virtual_ranges(self) -> None:
        headers = [(PT_LOAD, PF_R | PF_X, 0x1000, 0x1000, 0, 0x100, 0x1800, 0x1000),
                   (PT_LOAD, PF_R | PF_W, 0x2000, 0x2000, 0, 0x100, 0x100, 0x1000)]
        with self.assertRaises(ElfContractError):
            assess(make_elf(headers))

    def test_entry_must_be_executable(self) -> None:
        with self.assertRaises(ElfContractError):
            assess(make_elf(entry=0x3000))

    def test_program_header_count_is_bounded(self) -> None:
        headers = [(PT_LOAD, PF_R | PF_X, 0x1000, 0x1000, 0, 1, 1, 0x1000)] * 33
        with self.assertRaises(ElfContractError):
            assess(make_elf(headers))


if __name__ == "__main__":
    unittest.main(argv=[__file__])