#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Synthetic tests for the Phase-1.0AH dynamic/relocation contract."""

from __future__ import annotations

import argparse
import hashlib
from pathlib import Path
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
from phase10ah_dynamic_contract import *  # noqa: E402,F403


NEEDED = ("libSceLibcInternal.sprx", "libkernel_web.sprx")


def make_dynamic(relocations=None, needed=NEEDED, dynamic_terminated=True) -> bytes:
    strings = b"\0" + b"".join(name.encode("ascii") + b"\0" for name in needed)
    offsets = []
    position = 1
    for name in needed:
        offsets.append(position)
        position += len(name) + 1
    dynamic = b"".join(DYNAMIC_ENTRY.pack(DT_NEEDED, offset) for offset in offsets)
    if dynamic_terminated:
        dynamic += DYNAMIC_ENTRY.pack(DT_NULL, 0)
    if relocations is None:
        relocations = [(0x2100, R_X86_64_RELATIVE, 0, 0x100),
                       (0x2108, R_X86_64_GLOB_DAT, 1, 0)]
    rela = b"".join(RELA_ENTRY.pack(target, symbol << 32 | kind, addend)
                    for target, kind, symbol, addend in relocations)
    payload = bytearray(0x3400)
    ident = bytearray(16)
    ident[:7] = b"\x7fELF\x02\x01\x01"
    ELF_HEADER.pack_into(payload, 0, bytes(ident), ET_DYN, EM_X86_64, 1, 0x100,
                         64, 0x3000, 0, 64, 56, 2, 64, 4, 0)
    PROGRAM_HEADER.pack_into(payload, 64, PT_LOAD, PF_R | PF_X, 0, 0, 0,
                             0x1000, 0x1000, 0x1000)
    PROGRAM_HEADER.pack_into(payload, 120, PT_LOAD, PF_R | PF_W, 0x1000,
                             0x2000, 0, 0x1000, 0x2000, 0x1000)
    payload[0x1800:0x1800 + len(strings)] = strings
    payload[0x1900:0x1900 + len(dynamic)] = dynamic
    payload[0x1a00:0x1a00 + len(rela)] = rela
    sections = [
        (0, 0, 0, 0, 0, 0, 0, 0, 0, 0),
        (0, SHT_STRTAB, 0, 0, 0x1800, len(strings), 0, 0, 1, 0),
        (0, SHT_DYNAMIC, 0, 0, 0x1900, len(dynamic), 1, 0, 8, 16),
        (0, SHT_RELA, 0, 0, 0x1a00, len(rela), 0, 0, 8, 24),
    ]
    for index, section in enumerate(sections):
        SECTION_HEADER.pack_into(payload, 0x3000 + index * 64, *section)
    return bytes(payload)


def assess(payload: bytes, needed=NEEDED):
    return assess_dynamic(payload, hashlib.sha256(payload).hexdigest(), needed)


class DynamicContractTests(unittest.TestCase):
    def test_valid_split(self) -> None:
        result = assess(make_dynamic())
        self.assertEqual(result.relative_count, 1)
        self.assertEqual(result.glob_dat_count, 1)
        self.assertEqual(result.loader_applied_types, (R_X86_64_RELATIVE,))
        self.assertEqual(result.crt_applied_types, (R_X86_64_GLOB_DAT,))
        self.assertFalse(result.target_mapping_performed)

    def test_needed_inventory_is_exact_and_allowlisted(self) -> None:
        payload = make_dynamic()
        with self.assertRaises(DynamicContractError):
            assess(payload, tuple(reversed(NEEDED)))
        with self.assertRaises(DynamicContractError):
            assess(make_dynamic(needed=("evil.sprx",)), ("evil.sprx",))

    def test_dynamic_table_must_terminate(self) -> None:
        with self.assertRaises(DynamicContractError):
            assess(make_dynamic(dynamic_terminated=False))

    def test_unknown_relocation_type_is_rejected(self) -> None:
        with self.assertRaises(DynamicContractError):
            assess(make_dynamic([(0x2100, 7, 0, 0x100)]))

    def test_targets_must_be_aligned_and_writable(self) -> None:
        for target in (0x100, 0x2101, 0x5000):
            with self.subTest(target=target), self.assertRaises(DynamicContractError):
                assess(make_dynamic([(target, R_X86_64_RELATIVE, 0, 0x100)]))

    def test_relative_requires_zero_symbol_and_mapped_addend(self) -> None:
        for symbol, addend in ((1, 0x100), (0, -1), (0, 0x9000)):
            with self.subTest(), self.assertRaises(DynamicContractError):
                assess(make_dynamic([(0x2100, R_X86_64_RELATIVE, symbol, addend)]))

    def test_glob_dat_requires_symbol_and_zero_addend(self) -> None:
        for symbol, addend in ((0, 0), (1, 1)):
            relocs = [(0x2100, R_X86_64_RELATIVE, 0, 0x100),
                      (0x2108, R_X86_64_GLOB_DAT, symbol, addend)]
            with self.subTest(), self.assertRaises(DynamicContractError):
                assess(make_dynamic(relocs))

    def test_relative_closure_is_required(self) -> None:
        with self.assertRaises(DynamicContractError):
            assess(make_dynamic([(0x2108, R_X86_64_GLOB_DAT, 1, 0)]))

    def test_base_elf_hash_and_shape_remain_enforced(self) -> None:
        payload = make_dynamic()
        with self.assertRaises(DynamicContractError):
            assess_dynamic(payload, "0" * 64, NEEDED)


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