#!/usr/bin/env python3
"""用 Unicorn 离线执行 libweapon 的 VIMG 核心，作为纯 Python 移植的比对 oracle。

该工具只用于逆向阶段生成测试向量。生产代码不得依赖 Unicorn、ELF 样本或
本模块。运行示例：

    uv run --with pyelftools --with unicorn \
      python -m tools.libweapon_vimg_oracle --payload '{"probe":"vimg"}'
"""

from __future__ import annotations

import argparse
import struct
from pathlib import Path
from typing import Callable


DEFAULT_LIBRARY = Path(
    "out/p0_64_extract/lib/arm64-v8a/libweapon.2.2174a68a..so"
)
BASE_ENTRY = 0x1A1850
FULL_ENTRY = 0x19EA90

STACK_BASE = 0x70000000
STACK_SIZE = 0x400000
AUX_BASE = 0x71000000
AUX_SIZE = 0x10000
HEAP_BASE = 0x72000000
HEAP_SIZE = 0x1000000


class OracleError(RuntimeError):
    """离线 native oracle 执行失败。"""


class LibweaponVimgOracle:
    def __init__(self, library: Path = DEFAULT_LIBRARY) -> None:
        self.library = Path(library)

    def generate_base(self, payload: str) -> str:
        return self._run(BASE_ENTRY, payload)

    def generate_full(self, payload: str) -> str:
        return self._run(FULL_ENTRY, payload)

    def _run(
        self,
        entry: int,
        payload: str,
        code_observer: Callable[[object, int], None] | None = None,
    ) -> str:
        try:
            from elftools.elf.elffile import ELFFile
            from unicorn import UC_ARCH_ARM64, UC_HOOK_CODE, UC_MODE_LITTLE_ENDIAN, Uc
            from unicorn.arm64_const import (
                UC_ARM64_REG_SP,
                UC_ARM64_REG_TPIDR_EL0,
                UC_ARM64_REG_PC,
                UC_ARM64_REG_X0,
                UC_ARM64_REG_X1,
                UC_ARM64_REG_X2,
                UC_ARM64_REG_X3,
                UC_ARM64_REG_X4,
                UC_ARM64_REG_X30,
            )
        except ImportError as exc:
            raise OracleError(
                "需要通过 `uv run --with pyelftools --with unicorn` 运行 oracle"
            ) from exc

        if not self.library.is_file():
            raise OracleError(f"libweapon 样本不存在: {self.library}")

        with self.library.open("rb") as stream:
            elf = ELFFile(stream)
            loads = [segment for segment in elf.iter_segments() if segment["p_type"] == "PT_LOAD"]
            if not loads:
                raise OracleError("ELF 不含 PT_LOAD 段")

            page = 0x1000
            image_end = max(
                (int(segment["p_vaddr"]) + int(segment["p_memsz"]) + page - 1)
                & -page
                for segment in loads
            )
            emulator = Uc(UC_ARCH_ARM64, UC_MODE_LITTLE_ENDIAN)
            emulator.mem_map(0, image_end)
            for segment in loads:
                emulator.mem_write(int(segment["p_vaddr"]), segment.data())

            for section in elf.iter_sections():
                if not section.name.startswith(".rela"):
                    continue
                for relocation in section.iter_relocations():
                    if relocation.is_RELA():
                        emulator.mem_write(
                            int(relocation["r_offset"]),
                            struct.pack("<Q", int(relocation["r_addend"]) & 0xFFFFFFFFFFFFFFFF),
                        )

            rela_plt = elf.get_section_by_name(".rela.plt")
            plt = elf.get_section_by_name(".plt")
            if rela_plt is None or plt is None:
                raise OracleError("ELF 缺少 .plt/.rela.plt")
            dynsym = elf.get_section(rela_plt["sh_link"])
            plt_base = int(plt["sh_addr"]) + 0x20
            plt_names = {
                plt_base + index * 0x10: dynsym.get_symbol(rel["r_info_sym"]).name
                for index, rel in enumerate(rela_plt.iter_relocations())
            }

        emulator.mem_map(STACK_BASE, STACK_SIZE)
        emulator.mem_map(AUX_BASE, AUX_SIZE)
        emulator.mem_map(HEAP_BASE, HEAP_SIZE)

        tls = AUX_BASE + 0x1000
        input_address = AUX_BASE + 0x4000
        sentinel = AUX_BASE + 0x6000
        payload_bytes = payload.encode("utf-8")
        if b"\0" in payload_bytes:
            raise OracleError("payload 不得包含 NUL")
        if len(payload_bytes) + 1 > 0x2000:
            raise OracleError("oracle payload 超过 8191 字节")

        emulator.reg_write(UC_ARM64_REG_TPIDR_EL0, tls)
        emulator.mem_write(tls + 0x28, struct.pack("<Q", 0x123456789ABCDEF0))
        emulator.mem_write(input_address, payload_bytes + b"\0")
        emulator.reg_write(UC_ARM64_REG_SP, STACK_BASE + STACK_SIZE - 0x10000)
        emulator.reg_write(UC_ARM64_REG_X0, input_address)
        emulator.reg_write(UC_ARM64_REG_X30, sentinel)

        heap_next = HEAP_BASE + 0x1000

        def read_c_string(address: int, limit: int = 0x10000) -> bytes:
            if address == 0:
                return b""
            data = bytes(emulator.mem_read(address, limit))
            return data.split(b"\0", 1)[0]

        def return_from_call(value: int = 0) -> None:
            emulator.reg_write(UC_ARM64_REG_X0, value)
            emulator.reg_write(
                # Unicorn 在 hook 中修改 PC 即可跳过 PLT stub。
                UC_ARM64_REG_PC,
                emulator.reg_read(UC_ARM64_REG_X30),
            )

        def read_va_int(va_list: int) -> int:
            stack_pointer, gr_top, _vr_top, gr_offs, _vr_offs = struct.unpack(
                "<QQQii", bytes(emulator.mem_read(va_list, 32))
            )
            source = gr_top + gr_offs if gr_offs < 0 else stack_pointer
            return struct.unpack("<Q", bytes(emulator.mem_read(source, 8)))[0]

        def hook_code(_uc, address: int, _size: int, _user_data) -> None:
            nonlocal heap_next
            if code_observer is not None:
                code_observer(emulator, address)
            name = plt_names.get(address)
            if name is None:
                return

            x0 = emulator.reg_read(UC_ARM64_REG_X0)
            x1 = emulator.reg_read(UC_ARM64_REG_X1)
            x2 = emulator.reg_read(UC_ARM64_REG_X2)
            x3 = emulator.reg_read(UC_ARM64_REG_X3)
            x4 = emulator.reg_read(UC_ARM64_REG_X4)

            if name in {"malloc", "calloc", "realloc"}:
                size = x0 * x1 if name == "calloc" else (x1 if name == "realloc" else x0)
                result = (heap_next + 15) & -16
                heap_next = result + max(size, 16)
                if heap_next > HEAP_BASE + HEAP_SIZE:
                    raise OracleError("oracle heap 耗尽")
                emulator.mem_write(result, b"\0" * max(size, 1))
                return_from_call(result)
                return
            if name == "free":
                return_from_call()
                return
            if name in {"memcpy", "memmove"}:
                emulator.mem_write(x0, bytes(emulator.mem_read(x1, x2)))
                return_from_call(x0)
                return
            if name == "memset":
                emulator.mem_write(x0, bytes([x1 & 0xFF]) * x2)
                return_from_call(x0)
                return
            if name == "strlen":
                return_from_call(len(read_c_string(x0)))
                return
            if name == "__vsprintf_chk":
                format_string = read_c_string(x3).decode("ascii")
                value = read_va_int(x4)
                if format_string == "%08x ":
                    rendered = f"{value & 0xFFFFFFFF:08x} ".encode("ascii")
                elif format_string == "%02x":
                    rendered = f"{value & 0xFF:02x}".encode("ascii")
                else:
                    raise OracleError(f"未实现的 __vsprintf_chk 格式: {format_string!r}")
                emulator.mem_write(x0, rendered + b"\0")
                return_from_call(len(rendered))
                return
            if name.startswith("pthread_"):
                return_from_call()
                return
            if name == "__stack_chk_fail":
                raise OracleError("native 栈保护校验失败")
            raise OracleError(f"未实现的 PLT 调用: {name}")

        emulator.hook_add(UC_HOOK_CODE, hook_code)
        try:
            emulator.emu_start(entry, sentinel, count=2_000_000)
        except Exception as exc:
            if isinstance(exc, OracleError):
                raise
            raise OracleError(
                f"native 执行失败，PC={emulator.reg_read(UC_ARM64_REG_PC):#x}: {exc}"
            ) from exc

        result_address = emulator.reg_read(UC_ARM64_REG_X0)
        return read_c_string(result_address).decode("ascii")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--library", type=Path, default=DEFAULT_LIBRARY)
    parser.add_argument("--payload", required=True)
    parser.add_argument("--base-only", action="store_true")
    args = parser.parse_args()

    oracle = LibweaponVimgOracle(args.library)
    result = oracle.generate_base(args.payload) if args.base_only else oracle.generate_full(args.payload)
    print(result)
    return 0


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