#!/usr/bin/env python3
"""
Locate and apply the direct-panic mtepanichelper hook inside an Apple kernel
collection / Mach-O fileset.

This script is for the direct-panic helper design:

    XNU handle_kernel_tag_check_fault() fatal formatter block
        ADRL/MOV ... MSG_FMT
        MOV X0, MSG_FMT_REG
        BL  strlen(MSG_FMT)         <-- patch this instruction to BL helper
        ... original formatter code ...  <-- must not execute
        MOV X0, X23
        MOV X1, X19
        BL  panic_with_thread_kernel_state

The helper must be present in the same final KC/fileset before this script is
run. The helper entry is found by scanning the helper fileset entry for two
repeated magic qwords:

    0x2222D5D21111F5F2
    0x2222D5D21111F5F2

The branch target is magic+0x10, which should land on the BTI at the real hook
entry. This avoids relying on the custom __mtehook section name surviving KC
construction.

The patch source is the first BL after the MSG_FMT address load. On the provided
kernel build this is the strlen(MSG_FMT) call. The helper is entered with LR set
to the instruction after this original call site, scans forward from that LR for
"MOV X0, X23; MOV X1, X19; BL ...", builds its own message, calls the decoded
panic_with_thread_kernel_state() target itself, and never returns.

Place this script next to macho_fileset_branch_v3.py, or pass --branch-tool.

Examples:

    # Only locate / verify candidates
    python3 mte_apply_directpanic_hook.py locate KC_WITH_HELPER \
        --branch-tool ./macho_fileset_branch_v3.py

    # Patch candidate #1 and write a new KC
    python3 mte_apply_directpanic_hook.py patch KC_WITH_HELPER \
        --branch-tool ./macho_fileset_branch_v3.py \
        --output KC_WITH_HELPER_HOOKED

    # If multiple candidates are printed, select another one explicitly
    python3 mte_apply_directpanic_hook.py patch KC_WITH_HELPER \
        --candidate 2 \
        --output KC_WITH_HELPER_HOOKED
"""

from __future__ import annotations

import argparse
import importlib.util
import pathlib
import struct
import sys
from dataclasses import dataclass
from typing import Iterable, List, Optional, Sequence, Tuple

DEFAULT_KERNEL_ENTRY = "com.apple.kernel"
DEFAULT_HELPER_ENTRY = "com.binarygecko.mtepanichelper"
DEFAULT_MSG_FMT = b"Kernel tag check fault (expected tagged address: 0x%016llx)"
DEFAULT_MAGIC_QWORD = 0x2222D5D21111F5F2
DEFAULT_SCAN_INSNS = 256


@dataclass
class CodeRange:
    segname: str
    sectname: str
    va: int
    fileoff: int
    size: int

    def loc_expr(self, pc: int) -> str:
        return f"{self.segname}:{self.sectname}+0x{pc - self.va:x}"


@dataclass
class HookTarget:
    entry: str
    magic_va: int
    magic_fileoff: int
    magic_loc: str
    target_va: int
    target_loc: str


@dataclass
class PanicCallPattern:
    mov0_pc: int
    mov1_pc: int
    bl_pc: int
    bl_fileoff: int
    target_va: int
    loc: str


@dataclass
class Candidate:
    string_va: int
    string_fileoff: int
    string_loc: Optional[str]
    xref_pc: int
    xref_fileoff: int
    xref_loc: str
    msg_reg: int
    xref_kind: str
    patch_pc: int
    patch_fileoff: int
    patch_loc: str
    original_bl_target: int
    mov_x0_from_msg_reg: bool
    panic_pattern: Optional[PanicCallPattern]


def die(msg: str, code: int = 1) -> None:
    print(f"error: {msg}", file=sys.stderr)
    raise SystemExit(code)


def sign_extend(value: int, bits: int) -> int:
    sign = 1 << (bits - 1)
    return (value ^ sign) - sign


def read_u32(data: bytes, off: int) -> int:
    return struct.unpack_from("<I", data, off)[0]


def decode_adr(insn: int, pc: int) -> Optional[Tuple[int, int]]:
    # ADR Xd, imm
    if (insn & 0x9F000000) != 0x10000000:
        return None
    immlo = (insn >> 29) & 0x3
    immhi = (insn >> 5) & 0x7FFFF
    imm = sign_extend((immhi << 2) | immlo, 21)
    rd = insn & 0x1F
    return pc + imm, rd


def decode_adrp(insn: int, pc: int) -> Optional[Tuple[int, int]]:
    # ADRP Xd, imm
    if (insn & 0x9F000000) != 0x90000000:
        return None
    immlo = (insn >> 29) & 0x3
    immhi = (insn >> 5) & 0x7FFFF
    imm = sign_extend((immhi << 2) | immlo, 21) << 12
    rd = insn & 0x1F
    return (pc & ~0xFFF) + imm, rd


def decode_add_imm(insn: int) -> Optional[Tuple[int, int, int]]:
    # ADD Xd, Xn, #imm12{, LSL #12}; enough for ADRP+ADD address materialization.
    if (insn & 0x7F000000) != 0x11000000:
        return None
    shift = (insn >> 22) & 0x3
    if shift not in (0, 1):
        return None
    imm12 = (insn >> 10) & 0xFFF
    imm = imm12 << (12 if shift else 0)
    rn = (insn >> 5) & 0x1F
    rd = insn & 0x1F
    return rd, rn, imm


def decode_mov_reg(insn: int) -> Optional[Tuple[int, int]]:
    # MOV Xd, Xm == ORR Xd, XZR, Xm, 64-bit variant.
    if (insn & 0xFFE0FFE0) != 0xAA0003E0:
        return None
    dst = insn & 0x1F
    src = (insn >> 16) & 0x1F
    return dst, src


def decode_bl(insn: int, pc: int) -> Optional[int]:
    if (insn & 0xFC000000) != 0x94000000:
        return None
    imm26 = sign_extend(insn & 0x03FFFFFF, 26) << 2
    return pc + imm26


def load_branch_tool(branch_tool: pathlib.Path):
    branch_tool = branch_tool.resolve()
    if not branch_tool.exists():
        die(f"cannot find branch tool: {branch_tool}")
    spec = importlib.util.spec_from_file_location("macho_fileset_branch_v3", branch_tool)
    if spec is None or spec.loader is None:
        die(f"cannot import branch tool: {branch_tool}")
    mod = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = mod
    spec.loader.exec_module(mod)
    return mod


def default_branch_tool_path(script_path: pathlib.Path) -> pathlib.Path:
    return script_path.resolve().parent / "macho_fileset_branch_v3.py"


def fileoff_to_va(entry, fileoff: int) -> Optional[int]:
    base = 0 if entry.macho.fileoff_mode == "absolute" else entry.macho.root_offset
    for seg in entry.macho.segments:
        start = base + seg.fileoff
        end = start + seg.filesize
        if start <= fileoff < end:
            return seg.vmaddr + (fileoff - start)
    return None


def va_to_loc(entry, va: int) -> Optional[str]:
    sec = entry.macho.section_containing_va(va)
    if sec is not None:
        return f"{sec.segname}:{sec.sectname}+0x{va - sec.addr:x}"
    seg = entry.macho.segment_containing_va(va)
    if seg is not None:
        return f"{seg.segname}+0x{va - seg.vmaddr:x}"
    return None


def iter_code_ranges(entry) -> Iterable[CodeRange]:
    for sec in entry.macho.sections_by_index:
        if sec.size == 0:
            continue
        has_insns = (sec.flags & 0x80000000) or (sec.flags & 0x00000400)
        looks_like_text = sec.segname == "__TEXT_EXEC" and sec.sectname == "__text"
        if not (has_insns or looks_like_text):
            continue
        try:
            fileoff = entry.macho.va_to_fileoff(sec.addr)
        except Exception:
            continue
        yield CodeRange(sec.segname, sec.sectname, sec.addr, fileoff, sec.size)


def find_bytes_in_entry(data: bytes, entry, needle: bytes) -> List[Tuple[int, int]]:
    hits: List[Tuple[int, int]] = []
    start = 0
    while True:
        off = data.find(needle, start)
        if off < 0:
            break
        va = fileoff_to_va(entry, off)
        if va is not None:
            hits.append((va, off))
        start = off + 1
    return hits


def find_hook_target(data: bytes, helper_entry, helper_entry_name: str, magic_qword: int) -> HookTarget:
    magic = struct.pack("<Q", magic_qword) * 2
    hits = find_bytes_in_entry(data, helper_entry, magic)
    if not hits:
        die(
            f"did not find repeated hook magic 0x{magic_qword:016x} in {helper_entry_name}; "
            "is the direct-panic mtepanichelper present in this KC?"
        )
    if len(hits) > 1:
        print(f"warning: found {len(hits)} helper magic matches; using the first", file=sys.stderr)
    magic_va, magic_fileoff = hits[0]
    target_va = magic_va + len(magic)
    return HookTarget(
        entry=helper_entry_name,
        magic_va=magic_va,
        magic_fileoff=magic_fileoff,
        magic_loc=va_to_loc(helper_entry, magic_va) or f"+0x{magic_va - helper_entry.vmaddr:x}",
        target_va=target_va,
        target_loc=va_to_loc(helper_entry, target_va) or f"+0x{target_va - helper_entry.vmaddr:x}",
    )


def find_first_bl_after(data: bytes, cr: CodeRange, after_fileoff: int, max_insns: int) -> Optional[Tuple[int, int, int]]:
    start = after_fileoff + 4
    end = min(cr.fileoff + cr.size, start + max_insns * 4, len(data))
    off = start
    while off + 4 <= end:
        pc = cr.va + (off - cr.fileoff)
        target = decode_bl(read_u32(data, off), pc)
        if target is not None:
            return pc, off, target
        off += 4
    return None


def has_mov_x0_from_reg_before(data: bytes, from_fileoff: int, to_fileoff: int, reg: int) -> bool:
    off = from_fileoff + 4
    while off + 4 <= to_fileoff:
        mov = decode_mov_reg(read_u32(data, off))
        if mov == (0, reg):
            return True
        off += 4
    return False


def find_panic_call_pattern(data: bytes, entry, start_va: int, start_fileoff: int, max_insns: int) -> Optional[PanicCallPattern]:
    # Pattern expected by the direct-panic helper scanner:
    #   MOV X0, X23  (0xaa1703e0)
    #   MOV X1, X19  (0xaa1303e1)
    #   BL target
    mov_x0_x23 = 0xAA1703E0
    mov_x1_x19 = 0xAA1303E1
    end = min(len(data), start_fileoff + max_insns * 4)
    off = start_fileoff
    while off + 12 <= end:
        pc = fileoff_to_va(entry, off)
        if pc is None:
            off += 4
            continue
        insn0 = read_u32(data, off)
        insn1 = read_u32(data, off + 4)
        insn2 = read_u32(data, off + 8)
        target = decode_bl(insn2, pc + 8)
        if insn0 == mov_x0_x23 and insn1 == mov_x1_x19 and target is not None:
            return PanicCallPattern(
                mov0_pc=pc,
                mov1_pc=pc + 4,
                bl_pc=pc + 8,
                bl_fileoff=off + 8,
                target_va=target,
                loc=va_to_loc(entry, pc) or f"VA:0x{pc:x}",
            )
        off += 4
    return None


def find_candidates(
    data: bytes,
    kernel_entry,
    msg_fmt: bytes,
    max_adrp_window: int,
    max_first_bl_window: int,
    max_panic_scan_insns: int,
) -> List[Candidate]:
    strings = find_bytes_in_entry(data, kernel_entry, msg_fmt)
    candidates: List[Candidate] = []

    for string_va, string_fileoff in strings:
        string_page = string_va & ~0xFFF
        for cr in iter_code_ranges(kernel_entry):
            end = min(cr.fileoff + cr.size, len(data))
            off = cr.fileoff
            while off + 4 <= end:
                pc = cr.va + (off - cr.fileoff)
                insn = read_u32(data, off)

                adr = decode_adr(insn, pc)
                if adr is not None:
                    target, rd = adr
                    if target == string_va:
                        first = find_first_bl_after(data, cr, off, max_first_bl_window)
                        if first is not None:
                            patch_pc, patch_off, original_target = first
                            lr_after_patch = patch_pc + 4
                            panic_pattern = find_panic_call_pattern(
                                data, kernel_entry, lr_after_patch, patch_off + 4, max_panic_scan_insns
                            )
                            candidates.append(Candidate(
                                string_va=string_va,
                                string_fileoff=string_fileoff,
                                string_loc=va_to_loc(kernel_entry, string_va),
                                xref_pc=pc,
                                xref_fileoff=off,
                                xref_loc=cr.loc_expr(pc),
                                msg_reg=rd,
                                xref_kind="ADR",
                                patch_pc=patch_pc,
                                patch_fileoff=patch_off,
                                patch_loc=cr.loc_expr(patch_pc),
                                original_bl_target=original_target,
                                mov_x0_from_msg_reg=has_mov_x0_from_reg_before(data, off, patch_off, rd),
                                panic_pattern=panic_pattern,
                            ))

                adrp = decode_adrp(insn, pc)
                if adrp is not None:
                    page, rd = adrp
                    if page == string_page:
                        look = off + 4
                        limit = min(end, off + 4 + max_adrp_window * 4)
                        while look + 4 <= limit:
                            insn2 = read_u32(data, look)
                            add = decode_add_imm(insn2)
                            if add is not None:
                                add_rd, rn, imm = add
                                if rn == rd and page + imm == string_va:
                                    first = find_first_bl_after(data, cr, look, max_first_bl_window)
                                    if first is not None:
                                        patch_pc, patch_off, original_target = first
                                        lr_after_patch = patch_pc + 4
                                        panic_pattern = find_panic_call_pattern(
                                            data, kernel_entry, lr_after_patch, patch_off + 4, max_panic_scan_insns
                                        )
                                        candidates.append(Candidate(
                                            string_va=string_va,
                                            string_fileoff=string_fileoff,
                                            string_loc=va_to_loc(kernel_entry, string_va),
                                            xref_pc=pc,
                                            xref_fileoff=off,
                                            xref_loc=cr.loc_expr(pc),
                                            msg_reg=add_rd,
                                            xref_kind="ADRP+ADD",
                                            patch_pc=patch_pc,
                                            patch_fileoff=patch_off,
                                            patch_loc=cr.loc_expr(patch_pc),
                                            original_bl_target=original_target,
                                            mov_x0_from_msg_reg=has_mov_x0_from_reg_before(data, look, patch_off, add_rd),
                                            panic_pattern=panic_pattern,
                                        ))
                            look += 4
                off += 4

    # Sort with the most plausible first: MOV X0,msg_reg before first BL and panic pattern exists.
    candidates.sort(key=lambda c: (not c.mov_x0_from_msg_reg, c.panic_pattern is None, c.patch_pc))
    return candidates


def format_context(data: bytes, entry, center_fileoff: int, before: int = 8, after: int = 8) -> List[str]:
    lines: List[str] = []
    start = max(0, center_fileoff - before * 4)
    end = min(len(data), center_fileoff + (after + 1) * 4)
    off = start - (start & 3)
    while off + 4 <= end:
        va = fileoff_to_va(entry, off)
        word = read_u32(data, off)
        marker = "=>" if off == center_fileoff else "  "
        if va is None:
            lines.append(f"{marker} file+0x{off:x} {word:08x}")
        else:
            loc = va_to_loc(entry, va) or "?"
            bl = decode_bl(word, va)
            suffix = f" ; BL 0x{bl:x}" if bl is not None else ""
            lines.append(f"{marker} 0x{va:016x}  file+0x{off:x}  {word:08x}  {loc}{suffix}")
        off += 4
    return lines


def print_candidate(idx: int, c: Candidate, data: bytes, kernel_entry, context: int) -> None:
    print(f"Candidate #{idx}")
    print(f"  string VA:        0x{c.string_va:x}")
    print(f"  string loc:       {c.string_loc or '?'}")
    print(f"  xref kind:        {c.xref_kind}")
    print(f"  xref VA:          0x{c.xref_pc:x}")
    print(f"  xref loc:         {c.xref_loc}")
    print(f"  MSG_FMT register: x{c.msg_reg}")
    print(f"  MOV X0,MSG reg:   {'yes' if c.mov_x0_from_msg_reg else 'NO / not seen'}")
    print(f"  patch BL VA:      0x{c.patch_pc:x}")
    print(f"  patch BL loc:     {c.patch_loc}")
    print(f"  original BL tgt:  0x{c.original_bl_target:x}  ; should be strlen(MSG_FMT)")
    if c.panic_pattern is None:
        print("  helper scan test: FAILED to find MOV X0,X23; MOV X1,X19; BL ...")
    else:
        print("  helper scan test: ok")
        print(f"    pattern loc:    {c.panic_pattern.loc}")
        print(f"    panic BL VA:    0x{c.panic_pattern.bl_pc:x}")
        print(f"    panic target:   0x{c.panic_pattern.target_va:x}")
    print("  context around patch BL:")
    for line in format_context(data, kernel_entry, c.patch_fileoff, context, context):
        print(f"    {line}")
    if c.panic_pattern is not None:
        print("  context around panic call pattern:")
        for line in format_context(data, kernel_entry, c.panic_pattern.bl_fileoff, context, context):
            print(f"    {line}")
    print()


def make_patch_spec(kernel_entry: str, patch_loc: str, helper_entry: str, target_loc: str) -> str:
    return f"{kernel_entry} {patch_loc} {helper_entry} {target_loc}"


def apply_direct_patch(
    tool,
    outer,
    candidate: Candidate,
    helper_target: HookTarget,
    kernel_entry_name: str,
    output: pathlib.Path,
) -> None:
    spec_text = make_patch_spec(kernel_entry_name, candidate.patch_loc, helper_target.entry, helper_target.target_loc)
    patch = tool.parse_patch_spec(spec_text)
    applied = tool.apply_patch(outer, patch, True, {})
    output.write_bytes(outer.data)

    print(f"Wrote {output}")
    print()
    print("Applied patch:")
    print(f"  source entry:   {applied.spec.src_entry}")
    print(f"  source loc:     {applied.spec.src_loc.raw}")
    print(f"  source VA:      0x{applied.src_va:x}")
    print(f"  source fileoff: 0x{applied.src_fileoff:x}")
    print(f"  original insn:  0x{applied.original_insn:08x}")
    print(f"  target entry:   {applied.spec.dst_entry}")
    print(f"  target loc:     {applied.spec.dst_loc.raw}")
    print(f"  target VA:      0x{applied.dst_va:x}")
    print(f"  new insn:       0x{applied.insn:08x} (BL)")
    print()


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Locate/apply the direct-panic mtepanichelper branch hook")
    sub = p.add_subparsers(dest="cmd", required=True)

    def add_common(sp):
        sp.add_argument("kernel_collection", help="KC / Mach-O fileset that already contains mtepanichelper")
        sp.add_argument("--branch-tool", default=None, help="path to macho_fileset_branch_v3.py")
        sp.add_argument("--kernel-entry", default=DEFAULT_KERNEL_ENTRY)
        sp.add_argument("--helper-entry", default=DEFAULT_HELPER_ENTRY)
        sp.add_argument("--msg-fmt", default=DEFAULT_MSG_FMT.decode("ascii"))
        sp.add_argument("--magic", default=f"0x{DEFAULT_MAGIC_QWORD:016x}", help="helper marker qword repeated twice")
        sp.add_argument("--fileoff-mode", choices=("auto", "relative", "absolute"), default="auto")
        sp.add_argument("--max-adrp-window", type=int, default=8)
        sp.add_argument("--max-first-bl-window", type=int, default=16)
        sp.add_argument("--max-panic-scan-insns", type=int, default=DEFAULT_SCAN_INSNS)
        sp.add_argument("--context", type=int, default=6)

    locate = sub.add_parser("locate", help="locate hook target and candidate patch sites")
    add_common(locate)

    patch = sub.add_parser("patch", help="apply the hook patch and write a new KC")
    add_common(patch)
    patch.add_argument("--candidate", type=int, default=1, help="1-based candidate number to patch")
    patch.add_argument("--output", required=True, help="output path for patched KC")
    patch.add_argument("--allow-missing-panic-pattern", action="store_true", help="patch even if the helper scan pattern was not found")
    patch.add_argument("--allow-no-mov-x0-check", action="store_true", help="patch even if MOV X0,MSG_FMT-register was not seen before the first BL")

    return p


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    script_path = pathlib.Path(__file__)
    branch_tool_path = pathlib.Path(args.branch_tool) if args.branch_tool else default_branch_tool_path(script_path)
    tool = load_branch_tool(branch_tool_path)

    kc_path = pathlib.Path(args.kernel_collection)
    outer = tool.parse_outer_macho(kc_path, args.fileoff_mode)
    data = bytes(outer.data)

    try:
        kernel_entry = outer.get_entry(args.kernel_entry)
    except Exception as e:
        die(f"cannot find kernel entry {args.kernel_entry!r}: {e}")

    try:
        helper_entry = outer.get_entry(args.helper_entry)
    except Exception as e:
        die(f"cannot find helper entry {args.helper_entry!r}: {e}")

    helper_target = find_hook_target(data, helper_entry, args.helper_entry, int(args.magic, 0))
    print("Helper hook target:")
    print(f"  entry:      {helper_target.entry}")
    print(f"  magic VA:   0x{helper_target.magic_va:x}")
    print(f"  magic loc:  {helper_target.magic_loc}")
    print(f"  target VA:  0x{helper_target.target_va:x}")
    print(f"  target loc: {helper_target.target_loc}")
    print()

    candidates = find_candidates(
        data=data,
        kernel_entry=kernel_entry,
        msg_fmt=args.msg_fmt.encode("utf-8"),
        max_adrp_window=args.max_adrp_window,
        max_first_bl_window=args.max_first_bl_window,
        max_panic_scan_insns=args.max_panic_scan_insns,
    )

    if not candidates:
        die("no candidate hook site found; inspect the MSG_FMT xref manually")

    print(f"Found {len(candidates)} candidate hook site(s).")
    print()
    for i, c in enumerate(candidates, 1):
        print_candidate(i, c, data, kernel_entry, args.context)
        print("  patch command equivalent:")
        print(
            "    python3 macho_fileset_branch_v3.py patch "
            f"{str(kc_path)!r} --link --patch "
            f"'{args.kernel_entry} {c.patch_loc} {args.helper_entry} {helper_target.target_loc}' "
            "--output KC_HOOKED"
        )
        print()

    if args.cmd == "locate":
        print("Locate-only mode. No file was modified.")
        return 0

    if not (1 <= args.candidate <= len(candidates)):
        die(f"--candidate must be in range 1..{len(candidates)}")

    chosen = candidates[args.candidate - 1]
    if not chosen.mov_x0_from_msg_reg and not args.allow_no_mov_x0_check:
        die("chosen candidate did not show MOV X0,MSG_FMT-register before the first BL; pass --allow-no-mov-x0-check to override")
    if chosen.panic_pattern is None and not args.allow_missing_panic_pattern:
        die("chosen candidate did not contain the helper's panic-call scan pattern; pass --allow-missing-panic-pattern to override")

    # Reparse so the bytearray has not been accidentally modified by anything else.
    outer_for_patch = tool.parse_outer_macho(kc_path, args.fileoff_mode)
    apply_direct_patch(tool, outer_for_patch, chosen, helper_target, args.kernel_entry, pathlib.Path(args.output))

    print("Next checks in your disassembler/debugger:")
    print("  1. The patched instruction should be the first strlen(MSG_FMT) BL.")
    print("  2. The new instruction should be a BL to the helper magic+0x10 target.")
    print("  3. LR in the helper should equal patched_BL_address+4.")
    print("  4. The helper scanner should find MOV X0,X23; MOV X1,X19; BL panic_with_thread_kernel_state.")
    return 0


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