#!/usr/bin/env python3
"""
Auto-discover hook signatures from the Plex Media Server binary.

For each hook target, the script tries these strategies in order:
  1. Broad structural pattern (opcodes with displacements/immediates wildcarded)
  2. String-anchored discovery: find a key string, find the LEA referencing it,
     backtrack to the function prologue, auto-generate a pattern by disassembling
     the prologue and wildcarding all displacement/immediate operands via capstone.
  3. Relationship-based: search within another discovered function's body.

If all strategies fail, the build fails with diagnostics.

Usage:
    python3 discover_patterns.py <PMS binary> -o patterns_generated.h
"""

import sys
import os
import struct
import argparse

try:
    from capstone import Cs, CS_ARCH_X86, CS_MODE_64
except ImportError:
    print("ERROR: capstone not installed. Run: pip install capstone", file=sys.stderr)
    sys.exit(2)

# Capstone operand type constants
CS_OP_REG = 1
CS_OP_IMM = 2
CS_OP_MEM = 3

# Capstone x86 register IDs
X86_REG_RIP = 41


# ── ELF parsing ────────────────────────────────────────────────────────────

def parse_elf_segments(data):
    if data[:4] != b'\x7fELF':
        raise ValueError("Not an ELF file")
    e_phoff = struct.unpack('<Q', data[32:40])[0]
    e_phentsize = struct.unpack('<H', data[54:56])[0]
    e_phnum = struct.unpack('<H', data[56:58])[0]
    segs = []
    for i in range(e_phnum):
        off = e_phoff + i * e_phentsize
        if struct.unpack('<I', data[off:off+4])[0] != 1:
            continue
        segs.append((
            struct.unpack('<Q', data[off+8:off+16])[0],     # p_offset
            struct.unpack('<Q', data[off+16:off+24])[0],    # p_vaddr
            struct.unpack('<Q', data[off+32:off+40])[0],    # p_filesz
        ))
    return segs


def file_to_vaddr(segments, file_off):
    for p_offset, p_vaddr, p_filesz in segments:
        if p_offset <= file_off < p_offset + p_filesz:
            return file_off - p_offset + p_vaddr
    return file_off


def vaddr_to_file(segments, vaddr):
    for p_offset, p_vaddr, p_filesz in segments:
        if p_vaddr <= vaddr < p_vaddr + p_filesz:
            return vaddr - p_vaddr + p_offset
    return vaddr


# ── Byte pattern matching ──────────────────────────────────────────────────

def parse_pattern(p):
    out = []
    for tok in p.split():
        if tok in ('??', '?'):
            out.append(None)
        else:
            out.append(int(tok, 16))
    return out


def find_matches(data, pattern):
    matches = []
    plen = len(pattern)
    for i in range(len(data) - plen + 1):
        if all(pb is None or data[i+j] == pb for j, pb in enumerate(pattern)):
            matches.append(i)
    return matches


def pattern_to_str(pat):
    return ' '.join(f'{b:02X}' if b is not None else '?' for b in pat)


# ── Instruction analysis (capstone) ────────────────────────────────────────

LEGACY_PREFIXES = {0x26, 0x2e, 0x36, 0x3e, 0x64, 0x65, 0x66, 0x67, 0xf0, 0xf2, 0xf3}


def find_modrm_pos(raw, size):
    """Find the ModRM byte position in an x86-64 instruction's raw bytes."""
    pos = 0
    while pos < size and raw[pos] in LEGACY_PREFIXES:
        pos += 1
    if pos < size and 0x40 <= raw[pos] <= 0x4f:
        pos += 1  # REX
    if pos < size and raw[pos] == 0xc5:
        pos += 2  # VEX 2-byte
    elif pos < size and raw[pos] == 0xc4:
        pos += 3  # VEX 3-byte
    if pos < size:
        if raw[pos] == 0x0f:
            pos += 1
            if pos < size and raw[pos] in (0x38, 0x3a):
                pos += 1
        pos += 1  # opcode byte
    return pos if pos < size else -1


def wildcard_instruction(insn):
    """Return a list of (byte|None) for an instruction, wildcarding all
    displacement and immediate operands. None means wildcard."""
    raw = list(insn.bytes)
    size = insn.size
    wildcard = [False] * size

    has_rip_mem = False
    has_nonrip_mem_disp = False
    mem_disp = 0
    has_imm = False

    for op in insn.operands:
        if op.type == CS_OP_IMM:
            has_imm = True
        elif op.type == CS_OP_MEM:
            if op.mem.base == X86_REG_RIP:
                has_rip_mem = True
            elif op.mem.disp != 0 and op.mem.base != 0:
                has_nonrip_mem_disp = True
                mem_disp = op.mem.disp

    # Determine immediate total size and how many LOW bytes to wildcard.
    # imm_total: full immediate width. imm_wc: how many low bytes to
    # wildcard (high bytes kept for specificity, e.g. 0x00 for small frames).
    imm_total = 0
    imm_wc = 0
    if has_imm:
        mnem = insn.mnemonic
        if mnem in ('call', 'jmp') or mnem.startswith('j'):
            imm_total = 4 if size >= 5 else 1
            imm_wc = imm_total
        elif mnem in ('sub', 'add', 'cmp'):
            if size == 4:
                imm_total = 1; imm_wc = 1
            else:
                imm_total = 4; imm_wc = 2  # wildcard low 2, keep high 2 zeros
        elif mnem == 'mov':
            if size >= 10:
                imm_total = 8; imm_wc = 8
            elif size >= 7:
                imm_total = 4; imm_wc = 4
            else:
                imm_total = 1; imm_wc = 1
        elif mnem == 'push':
            imm_total = 1 if size == 2 else 4; imm_wc = imm_total
        elif mnem == 'test':
            imm_total = 1 if size <= 4 else 4; imm_wc = imm_total
        else:
            imm_total = min(4, size - 1); imm_wc = imm_total

    # Determine displacement size
    disp_size = 0
    if has_rip_mem:
        disp_size = 4
    elif has_nonrip_mem_disp:
        modrm_pos = find_modrm_pos(raw, size)
        if modrm_pos >= 0:
            mod_field = (raw[modrm_pos] >> 6) & 3
            if mod_field == 1:
                disp_size = 1
            elif mod_field == 2:
                disp_size = 4
        if disp_size == 0:
            disp_size = 1 if -128 <= mem_disp <= 127 else 4

    # Wildcard displacement bytes (immediately before the immediate field)
    if disp_size > 0:
        disp_start = size - imm_total - disp_size
        for i in range(max(0, disp_start), min(size, disp_start + disp_size)):
            wildcard[i] = True

    # Wildcard the LOW imm_wc bytes of the immediate (keep high bytes)
    if imm_total > 0:
        imm_start = size - imm_total
        for i in range(max(0, imm_start), min(size, imm_start + imm_wc)):
            wildcard[i] = True

    return [(raw[i] if not wildcard[i] else None) for i in range(size)]


def auto_wildcard_pattern(data, segments, func_file_off, length):
    """Disassemble a function and generate a byte pattern with all
    displacement/immediate operands auto-wildcarded."""
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    md.detail = True
    vaddr = file_to_vaddr(segments, func_file_off)
    code = data[func_file_off:func_file_off + length + 32]

    pattern = []
    bytes_consumed = 0

    for insn in md.disasm(code, vaddr):
        if bytes_consumed >= length:
            break
        wc_bytes = wildcard_instruction(insn)
        for b in wc_bytes:
            if bytes_consumed >= length:
                break
            pattern.append(b)
            bytes_consumed += 1

    while len(pattern) < length:
        pattern.append(None)

    return pattern[:length]


# ── Function discovery ─────────────────────────────────────────────────────

def find_string_in_binary(data, search_string):
    """Find all occurrences of a null-terminated string in the binary."""
    needle = search_string.encode() + b'\x00'
    results = []
    start = 0
    while True:
        idx = data.find(needle, start)
        if idx == -1:
            break
        results.append(idx)
        start = idx + 1
    return results


def find_lea_refs_to_string(data, segments, string_file_off):
    """Find all LEA reg,[rip+disp32] instructions that reference a string."""
    results = []
    for p_offset, p_vaddr, p_filesz in segments:
        end = min(p_offset + p_filesz, len(data) - 7)
        for i in range(p_offset, end):
            if data[i] in (0x48, 0x4c) and data[i+1] == 0x8D:
                modrm = data[i+2]
                if (modrm & 0xC7) == 0x05:
                    disp = struct.unpack('<i', data[i+3:i+7])[0]
                    insn_vaddr = file_to_vaddr(segments, i)
                    target_vaddr = insn_vaddr + 7 + disp
                    target_file = vaddr_to_file(segments, target_vaddr)
                    if target_file == string_file_off:
                        reg = (modrm >> 3) & 7
                        if data[i] == 0x4c:
                            reg += 8
                        results.append((i, reg))
    return results


def find_func_start_backwards(data, file_off, max_scan=16384):
    """Scan backwards for a function prologue (55 48 89 E5)."""
    for j in range(file_off, max(0, file_off - max_scan), -1):
        if data[j:j+4] == b'\x55\x48\x89\xe5':
            return j
    return None


def find_func_end(data, segments, func_file_off, max_insns=5000):
    """Find the end of a function by scanning for ret/int3 after the prologue."""
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    vaddr = file_to_vaddr(segments, func_file_off)
    code = data[func_file_off:func_file_off + 16384]
    count = 0
    for insn in md.disasm(code, vaddr):
        count += 1
        if count > max_insns:
            break
        if insn.mnemonic in ('ret', 'repret', 'ud2'):
            return func_file_off + insn.address - vaddr + insn.size
    return func_file_off + 8192


def string_anchored_discovery(data, segments, search_string, pattern_length,
                               expected_func_start_prefix=None):
    """Discover a function by finding a string reference, then backtracking
    to the function prologue. Auto-generates a pattern with wildcarded
    displacement/immediate operands."""
    string_locations = find_string_in_binary(data, search_string)
    if not string_locations:
        return None, f"string '{search_string}' not found in binary"

    for string_off in string_locations:
        lea_refs = find_lea_refs_to_string(data, segments, string_off)
        for lea_off, reg in lea_refs:
            func_start = find_func_start_backwards(data, lea_off)
            if not func_start:
                continue
            if expected_func_start_prefix:
                prefix = data[func_start:func_start + len(expected_func_start_prefix)]
                if prefix != expected_func_start_prefix:
                    continue
            pattern = auto_wildcard_pattern(data, segments, func_start, pattern_length)
            matches = find_matches(data, pattern)
            if matches:
                return pattern, f"found via string '{search_string}' -> LEA at 0x{lea_off:08x} -> func at 0x{func_start:08x} ({len(matches)} match(es))"

    return None, f"string '{search_string}' found but no enclosing function prologue"


def relationship_based_discovery(data, segments, ref_func_off, search_pattern_str,
                                  search_range=8192):
    """Search within a function's body for a structural pattern."""
    pat = parse_pattern(search_pattern_str)
    end = min(len(data), ref_func_off + search_range)
    matches = []
    for i in range(ref_func_off, end - len(pat) + 1):
        if all(pb is None or data[i+j] == pb for j, pb in enumerate(pat)):
            matches.append(i)
    if matches:
        return pat, f"found {len(matches)} match(es) within function body"
    return None, "pattern not found within function body"


# ── Hook target definitions ────────────────────────────────────────────────

TARGETS = [
    {
        "name": "PREF_GETTER",
        "broad_pattern": "55 48 89 E5 41 57 41 56 53 48 83 EC ? 48 89 F3 49 89 FE 0F B6 46 17 48 89 F1 84 C0",
        "string_anchor": None,
        "relates_to": None,
        "expected_matches": (1, 3),
        "required": True,
        "length": 26,
    },
    {
        "name": "BITSET_REF",
        "broad_pattern": "48 8D 0D ? ? ? ? 48 8B 94 05 ? ? ? ? 48 87 14 08",
        "string_anchor": None,
        "relates_to": ("BS_INIT", "48 8D 0D ? ? ? ? 48 8B 94 05 ? ? ? ? 48 87 14 08"),
        "expected_matches": (1, 4),
        "required": True,
        "length": 18,
    },
    {
        "name": "BS_INIT",
        "broad_pattern": "55 48 89 E5 41 57 41 56 41 55 41 54 53 48 81 EC ? ? 00 00 49 89 FE 48 8D 9D ? ? ? ? 48 89 DF E8 ? ? ? ? 48 8B 1B 48 85 DB",
        "string_anchor": "//feature",
        "string_prefix": b'\x55\x48\x89\xe5\x41\x57\x41\x56\x41\x55\x41\x54\x53',
        "relates_to": None,
        "expected_matches": (1, 1),
        "required": True,
        "length": 44,
    },
    {
        "name": "LEGACY_USF",
        "broad_pattern": "55 48 89 E5 48 8B 07 48 85 C0 74 09",
        "string_anchor": None,
        "relates_to": None,
        "expected_matches": (1, 5),
        "required": False,
        "length": 12,
    },
    {
        "name": "LEGACY_MF",
        "broad_pattern": "55 48 89 E5 41 57 41 56 53 48 83 EC ? 49 89 F7 4C 8D 77",
        "string_anchor": None,
        "relates_to": None,
        "expected_matches": (1, 3),
        "required": False,
        "length": 18,
    },
]


def generate_header(patterns, output_path):
    with open(output_path, 'w') as f:
        f.write("#pragma once\n")
        f.write("// Auto-generated by discover_patterns.py — DO NOT EDIT.\n")
        f.write("// Hook signatures discovered from the PMS binary at build time.\n\n")
        for name, pattern_str, match_count, method in patterns:
            f.write(f'// {name}: {match_count} match(es) — {method}\n')
            f.write(f'static const char* PATTERN_{name} = "{pattern_str}";\n\n')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('pms_path')
    parser.add_argument('-o', '--output', default='patterns_generated.h')
    args = parser.parse_args()

    with open(args.pms_path, 'rb') as f:
        data = f.read()
    segments = parse_elf_segments(data)
    print(f"Loaded {args.pms_path} ({len(data)} bytes, {len(segments)} LOAD segments)")

    results = []
    discovered_func_offsets = {}
    all_ok = True

    # Pass 1: discover targets via broad pattern or string-anchored (independent)
    # Pass 2: discover relationship-dependent targets using results from pass 1
    pending = []
    for target in TARGETS:
        name = target['name']
        broad = parse_pattern(target['broad_pattern'])
        lo, hi = target['expected_matches']

        # Strategy 1: broad structural pattern
        matches = find_matches(data, broad)
        if lo <= len(matches) <= hi:
            print(f"  [OK]   {name}: broad pattern ({len(matches)} matches)")
            for m in matches[:3]:
                print(f"         0x{m:08x}: {data[m:m+min(len(broad)+8,32)].hex(' ')}")
            results.append((name, pattern_to_str(broad), len(matches), f"broad pattern ({len(matches)} matches)"))
            discovered_func_offsets[name] = matches[0] if matches else None
            continue

        # Strategy 2: string-anchored discovery
        if target.get('string_anchor'):
            pattern, detail = string_anchored_discovery(
                data, segments, target['string_anchor'],
                target['length'], target.get('string_prefix'))
            if pattern:
                matches = find_matches(data, pattern)
                if lo <= len(matches) <= hi:
                    print(f"  [OK]   {name}: string-anchored ({len(matches)} matches)")
                    print(f"         {detail}")
                    for m in matches[:3]:
                        print(f"         0x{m:08x}: {data[m:m+min(len(pattern)+8,32)].hex(' ')}")
                    results.append((name, pattern_to_str(pattern), len(matches), f"string-anchored: {detail}"))
                    discovered_func_offsets[name] = matches[0] if matches else None
                    continue

        # Defer relationship-based to pass 2
        if target.get('relates_to'):
            pending.append(target)
        elif target['required']:
            print(f"  [FAIL] {name}: all strategies failed")
            print(f"         Broad pattern: {pattern_to_str(broad)}")
            if target.get('string_anchor'):
                print(f"         String anchor: '{target['string_anchor']}'")
            all_ok = False
            results.append((name, pattern_to_str(broad), 0, "FAILED"))
        else:
            print(f"  [SKIP] {name}: not found (optional)")
            results.append((name, pattern_to_str(broad), 0, "skipped (optional)"))

    # Pass 2: relationship-based discovery (depends on pass 1 results)
    for target in pending:
        name = target['name']
        broad = parse_pattern(target['broad_pattern'])
        lo, hi = target['expected_matches']
        ref_name, rel_pattern = target['relates_to']
        ref_off = discovered_func_offsets.get(ref_name)

        if ref_off:
            # Search within the referenced function's body
            func_end = find_func_end(data, segments, ref_off)
            pattern, detail = relationship_based_discovery(
                data, segments, ref_off, rel_pattern, search_range=func_end - ref_off + 256)
            if pattern:
                matches = find_matches(data, pattern)
                if lo <= len(matches) <= hi:
                    print(f"  [OK]   {name}: relationship ({ref_name}) ({len(matches)} matches)")
                    print(f"         {detail}")
                    for m in matches[:3]:
                        print(f"         0x{m:08x}: {data[m:m+min(len(pattern)+8,32)].hex(' ')}")
                    results.append((name, pattern_to_str(pattern), len(matches), f"relationship ({ref_name}): {detail}"))
                    discovered_func_offsets[name] = matches[0] if matches else None
                    continue

        # Relationship failed — try broad as last resort
        matches = find_matches(data, broad)
        if lo <= len(matches) <= hi:
            print(f"  [OK]   {name}: broad pattern ({len(matches)} matches) [fallback]")
            results.append((name, pattern_to_str(broad), len(matches), f"broad pattern fallback ({len(matches)} matches)"))
            discovered_func_offsets[name] = matches[0] if matches else None
            continue

        if target['required']:
            print(f"  [FAIL] {name}: all strategies failed (broad + string + relationship)")
            all_ok = False
            results.append((name, pattern_to_str(broad), 0, "FAILED"))
        else:
            print(f"  [SKIP] {name}: not found (optional)")
            results.append((name, pattern_to_str(broad), 0, "skipped (optional)"))

    if all_ok:
        generate_header(results, args.output)
        print(f"\nAll required patterns discovered. Header written to {args.output}")
        sys.exit(0)
    else:
        print("\nFAILED: one or more required patterns not found.", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()