#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""libweapon.so 纯静态 OLLVM/kcode-guard 初筛工具。

用途：
- 不运行目标 so，不依赖 Frida / 抓包 / 真机。
- 用 ELF 元数据 + Capstone 反汇编做轻量画像：
  1. 节区、符号、导入、熵；
  2. JNI_OnLoad 周边反汇编；
  3. br/blr/adrp/movk 等 OLLVM/间接调度指标；
  4. .data/.data.rel.ro 中连续指向 .text 的指针表候选；
  5. 明文字符串及近似 ADRP 直连引用。

运行示例：
  uv run --with capstone --with pyelftools python -m tools.analyze_libweapon_ollvm \
    out/p0_64_extract/lib/arm64-v8a/libweapon.2.2174a68a..so \
    --json-out out/libweapon_ollvm_report.json \
    --md-out out/libweapon_ollvm_report.md
"""

from __future__ import annotations

import argparse
import collections
import hashlib
import json
import math
import re
import struct
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Iterable

from capstone import Cs, CS_ARCH_ARM64, CS_MODE_ARM
from elftools.elf.elffile import ELFFile


TARGET_STRINGS = [
    "VIMG_",
    "a_y_q_z",
    "gifshow",
    "wcfg",
    "putString",
    "getString",
    "getSharedPreferences",
    "RegisterNatives",
    "JNI_OnLoad",
    "KsBridge",
    "Engine",
    "/NP/vmp",
]


@dataclass(frozen=True)
class SectionInfo:
    name: str
    addr: int
    offset: int
    size: int
    flags: str
    entropy: float


@dataclass(frozen=True)
class SymbolInfo:
    name: str
    value: int
    size: int
    section: str
    bind: str
    typ: str


@dataclass(frozen=True)
class StringHit:
    pattern: str
    offset: int
    vaddr: int
    text: str
    direct_adrp_refs: list[int] = field(default_factory=list)


@dataclass(frozen=True)
class PointerRun:
    section: str
    file_offset: int
    vaddr: int
    count: int
    targets: list[int]


@dataclass(frozen=True)
class DisasmLine:
    address: int
    mnemonic: str
    op_str: str


@dataclass(frozen=True)
class PltXref:
    name: str
    plt: int
    calls: list[int]


@dataclass(frozen=True)
class OllvmMetrics:
    text_instruction_count: int
    branch_count: int
    indirect_branch_count: int
    call_count: int
    indirect_call_count: int
    ret_count: int
    adrp_count: int
    movk_count: int
    csel_family_count: int
    cond_branch_count: int
    suspicious_dispatch_windows: list[dict[str, Any]]
    mnemonic_top: list[tuple[str, int]]


@dataclass(frozen=True)
class Report:
    path: str
    size: int
    sha256: str
    elf_machine: str
    entry: int
    sections: list[SectionInfo]
    defined_function_count: int
    imported_function_count: int
    exports_interesting: list[SymbolInfo]
    imports_interesting: list[SymbolInfo]
    string_hits: list[StringHit]
    pointer_runs: list[PointerRun]
    plt_xrefs_interesting: list[PltXref]
    jni_onload: SymbolInfo | None
    jni_onload_disasm: list[DisasmLine]
    ollvm_metrics: OllvmMetrics


def _entropy(data: bytes) -> float:
    if not data:
        return 0.0
    counts = collections.Counter(data)
    total = len(data)
    return -sum((n / total) * math.log2(n / total) for n in counts.values())


def _section_flags(sec) -> str:
    flags = int(sec["sh_flags"])
    out = []
    if flags & 0x1:
        out.append("W")
    if flags & 0x2:
        out.append("A")
    if flags & 0x4:
        out.append("X")
    return "".join(out) or "-"


class ElfView:
    def __init__(self, path: Path):
        self.path = path
        self.data = path.read_bytes()
        self.elf = ELFFile(path.open("rb"))
        self.sections = [sec for sec in self.elf.iter_sections()]
        self.load_segments = [seg for seg in self.elf.iter_segments() if seg["p_type"] == "PT_LOAD"]
        self.rela_addends = self._load_rela_addends()

    def _load_rela_addends(self) -> dict[int, int]:
        """返回 r_offset -> r_addend，用于识别重定位写入的调度表。"""

        out: dict[int, int] = {}
        for sec in self.sections:
            if not sec.name.startswith(".rela"):
                continue
            try:
                for rel in sec.iter_relocations():
                    if rel.is_RELA():
                        out[int(rel["r_offset"])] = int(rel["r_addend"])
            except Exception:
                continue
        return out

    def section_by_name(self, name: str):
        return self.elf.get_section_by_name(name)

    def vaddr_to_offset(self, vaddr: int) -> int | None:
        for seg in self.load_segments:
            start = int(seg["p_vaddr"])
            end = start + int(seg["p_filesz"])
            if start <= vaddr < end:
                return int(seg["p_offset"]) + (vaddr - start)
        return None

    def offset_to_vaddr(self, offset: int) -> int | None:
        for seg in self.load_segments:
            start = int(seg["p_offset"])
            end = start + int(seg["p_filesz"])
            if start <= offset < end:
                return int(seg["p_vaddr"]) + (offset - start)
        return None

    def bytes_at_vaddr(self, vaddr: int, size: int) -> bytes:
        off = self.vaddr_to_offset(vaddr)
        if off is None:
            return b""
        return self.data[off : off + size]


def _iter_symbols(view: ElfView) -> Iterable[SymbolInfo]:
    for sec_name in (".dynsym", ".symtab"):
        sec = view.section_by_name(sec_name)
        if sec is None:
            continue
        for sym in sec.iter_symbols():
            name = sym.name
            if not name:
                continue
            info = sym["st_info"]
            yield SymbolInfo(
                name=name,
                value=int(sym["st_value"]),
                size=int(sym["st_size"]),
                section="UND" if sym["st_shndx"] == "SHN_UNDEF" else sec_name,
                bind=str(info["bind"]),
                typ=str(info["type"]),
            )


def _extract_ascii_strings(data: bytes, min_len: int = 4) -> list[tuple[int, str]]:
    regex = rb"[ -~]{" + str(min_len).encode("ascii") + rb",}"
    return [(m.start(), m.group(0).decode("ascii", errors="replace")) for m in re.finditer(regex, data)]


def _disasm_section(view: ElfView, name: str) -> list[Any]:
    sec = view.section_by_name(name)
    if sec is None:
        return []
    md = Cs(CS_ARCH_ARM64, CS_MODE_ARM)
    md.detail = True
    md.skipdata = True
    return [ins for ins in md.disasm(sec.data(), int(sec["sh_addr"])) if ins.mnemonic != ".byte"]


def _disasm_window(view: ElfView, vaddr: int, size: int) -> list[DisasmLine]:
    md = Cs(CS_ARCH_ARM64, CS_MODE_ARM)
    md.detail = False
    lines = []
    for ins in md.disasm(view.bytes_at_vaddr(vaddr, size), vaddr):
        lines.append(DisasmLine(address=int(ins.address), mnemonic=ins.mnemonic, op_str=ins.op_str))
    return lines


def _adrp_ref_map(text_insns: list[Any]) -> dict[int, list[int]]:
    """近似收集 ADRP 目标页 -> 指令地址。

    Capstone 对 ARM64 ADRP 的第二操作数通常给出归一后的页地址。这里用页粒度
    做粗筛，适合回答“明文字符串是否有直连 ADRP 引用”。
    """

    refs: dict[int, list[int]] = collections.defaultdict(list)
    for ins in text_insns:
        if ins.mnemonic != "adrp":
            continue
        try:
            if len(ins.operands) >= 2 and hasattr(ins.operands[1], "imm"):
                refs[int(ins.operands[1].imm) & ~0xFFF].append(int(ins.address))
        except Exception:
            continue
    return refs


def _string_hits(view: ElfView, text_insns: list[Any]) -> list[StringHit]:
    refs_by_page = _adrp_ref_map(text_insns)
    hits: list[StringHit] = []
    for off, text in _extract_ascii_strings(view.data):
        matched = [pat for pat in TARGET_STRINGS if pat.lower() in text.lower()]
        if not matched:
            continue
        vaddr = view.offset_to_vaddr(off) or 0
        page = vaddr & ~0xFFF
        refs = refs_by_page.get(page, [])
        for pat in matched:
            hits.append(
                StringHit(
                    pattern=pat,
                    offset=off,
                    vaddr=vaddr,
                    text=text[:240],
                    direct_adrp_refs=refs[:32],
                )
            )
    return hits


def _pointer_runs(view: ElfView, min_run: int = 4) -> list[PointerRun]:
    text = view.section_by_name(".text")
    if text is None:
        return []
    text_start = int(text["sh_addr"])
    text_end = text_start + int(text["sh_size"])
    candidates: list[PointerRun] = []
    # kcode-guard 在 JNI_OnLoad 里使用：
    #   ldr x8, [table, idx, uxtw #3]
    #   add x8, x8, #BASE_CONST
    #   br x8
    # 表项不是裸指针，而是 signed/unsigned 64-bit delta。这里同时尝试直连
    # 指针和常见 delta base，先覆盖当前样本的 0x7c585548。
    delta_bases = [0, 0x7C585548]
    for sec_name in (".data.rel.ro", ".data", ".bss.rel.ro", ".got", ".got.plt"):
        sec = view.section_by_name(sec_name)
        if sec is None or int(sec["sh_size"]) < 8:
            continue
        data = sec.data()
        run_targets: list[int] = []
        run_start_idx = 0

        def flush(idx: int) -> None:
            nonlocal run_targets, run_start_idx
            if len(run_targets) >= min_run:
                file_off = int(sec["sh_offset"]) + run_start_idx * 8
                vaddr = int(sec["sh_addr"]) + run_start_idx * 8
                candidates.append(
                    PointerRun(
                        section=sec_name,
                        file_offset=file_off,
                        vaddr=vaddr,
                        count=len(run_targets),
                        targets=run_targets[:64],
                    )
                )
            run_targets = []

        for idx in range(0, len(data) // 8):
            raw = struct.unpack_from("<Q", data, idx * 8)[0]
            signed_raw = struct.unpack_from("<q", data, idx * 8)[0]
            slot_vaddr = int(sec["sh_addr"]) + idx * 8
            rela_addend = view.rela_addends.get(slot_vaddr)
            matched = None
            for base in delta_bases:
                values = [raw + base, signed_raw + base]
                if rela_addend is not None:
                    values.append(rela_addend + base)
                for value in values:
                    # ARM64 指令地址 4 字节对齐；br 目标必须落在可执行节。
                    if text_start <= value < text_end and value % 4 == 0:
                        matched = value
                        break
                if matched is not None:
                    break
            if matched is not None:
                if not run_targets:
                    run_start_idx = idx
                run_targets.append(matched)
            else:
                flush(idx)
        flush(len(data) // 8)
    candidates.sort(key=lambda r: r.count, reverse=True)
    return candidates


def _ollvm_metrics(text_insns: list[Any]) -> OllvmMetrics:
    mnemonic_counter = collections.Counter(ins.mnemonic for ins in text_insns)
    branch_mnemonics = {"b", "b.eq", "b.ne", "b.cs", "b.hs", "b.cc", "b.lo", "b.mi", "b.pl", "b.vs", "b.vc", "b.hi", "b.ls", "b.ge", "b.lt", "b.gt", "b.le", "cbz", "cbnz", "tbz", "tbnz"}
    cond_branch_count = sum(1 for ins in text_insns if ins.mnemonic in branch_mnemonics and ins.mnemonic != "b")
    indirect = [ins for ins in text_insns if ins.mnemonic == "br"]
    indirect_call = [ins for ins in text_insns if ins.mnemonic == "blr"]
    branches = [ins for ins in text_insns if ins.mnemonic.startswith("b") or ins.mnemonic in {"cbz", "cbnz", "tbz", "tbnz"}]
    calls = [ins for ins in text_insns if ins.mnemonic in {"bl", "blr"}]
    csel_family = [ins for ins in text_insns if ins.mnemonic in {"csel", "csinc", "csinv", "csneg", "cset", "csetm"}]

    suspicious: list[dict[str, Any]] = []
    for i, ins in enumerate(text_insns):
        if ins.mnemonic != "br":
            continue
        window = text_insns[max(0, i - 8) : i + 1]
        if any(w.mnemonic == "ldr" and "x8" in w.op_str for w in window) or any(w.mnemonic == "movk" for w in window):
            suspicious.append(
                {
                    "br": int(ins.address),
                    "window": [
                        {"address": int(w.address), "mnemonic": w.mnemonic, "op_str": w.op_str}
                        for w in window
                    ],
                }
            )
    return OllvmMetrics(
        text_instruction_count=len(text_insns),
        branch_count=len(branches),
        indirect_branch_count=len(indirect),
        call_count=len(calls),
        indirect_call_count=len(indirect_call),
        ret_count=mnemonic_counter.get("ret", 0),
        adrp_count=mnemonic_counter.get("adrp", 0),
        movk_count=mnemonic_counter.get("movk", 0),
        csel_family_count=len(csel_family),
        cond_branch_count=cond_branch_count,
        suspicious_dispatch_windows=suspicious[:80],
        mnemonic_top=mnemonic_counter.most_common(40),
    )


def _plt_symbol_map(view: ElfView) -> dict[int, str]:
    """AArch64 PLT 地址 -> 导入符号名。

    当前样本 `.plt` 大小满足：PLT0(0x20) + N * 0x10。
    `.rela.plt` 的第 i 项通常对应 `.plt + 0x20 + i*0x10`。
    """

    plt = view.section_by_name(".plt")
    rela_plt = view.section_by_name(".rela.plt")
    dynsym = view.section_by_name(".dynsym")
    if plt is None or rela_plt is None or dynsym is None:
        return {}
    out: dict[int, str] = {}
    base = int(plt["sh_addr"]) + 0x20
    for idx, rel in enumerate(rela_plt.iter_relocations()):
        sym_idx = rel["r_info_sym"]
        sym = dynsym.get_symbol(sym_idx)
        if sym and sym.name:
            out[base + idx * 0x10] = sym.name
    return out


def _plt_xrefs(view: ElfView, text_insns: list[Any]) -> list[PltXref]:
    plt_map = _plt_symbol_map(view)
    calls_by_plt: dict[int, list[int]] = collections.defaultdict(list)
    for ins in text_insns:
        if ins.mnemonic != "bl":
            continue
        try:
            if ins.operands and hasattr(ins.operands[0], "imm"):
                target = int(ins.operands[0].imm)
                if target in plt_map:
                    calls_by_plt[target].append(int(ins.address))
        except Exception:
            continue
    interesting = re.compile(
        r"mprotect|mmap|munmap|dlopen|dlsym|deflate|inflate|open|fopen|read|"
        r"clock|gettimeofday|system_property|socket|connect|send|recv",
        re.I,
    )
    rows = [
        PltXref(name=plt_map[addr], plt=addr, calls=calls[:120])
        for addr, calls in sorted(calls_by_plt.items(), key=lambda item: plt_map[item[0]])
        if interesting.search(plt_map[addr]) or calls
    ]
    return rows


def analyze(path: Path) -> Report:
    view = ElfView(path)
    text_insns = _disasm_section(view, ".text")
    all_symbols = list(_iter_symbols(view))
    defined_funcs = [s for s in all_symbols if s.typ == "STT_FUNC" and s.section != "UND"]
    imported_funcs = [s for s in all_symbols if s.typ == "STT_FUNC" and s.section == "UND"]
    interesting_re = re.compile(r"JNI|Java_|RegisterNatives|KsBridge|Engine|SharedPreferences|putString|getString|dlopen|dlsym|mmap|mprotect|deflate|inflate", re.I)
    jni = next((s for s in defined_funcs if s.name == "JNI_OnLoad"), None)
    jni_disasm = _disasm_window(view, jni.value, min(max(jni.size, 0x240), 0x800)) if jni else []

    sections = []
    for sec in view.sections:
        name = sec.name
        if not name:
            continue
        data = sec.data() if sec["sh_type"] != "SHT_NOBITS" else b""
        sections.append(
            SectionInfo(
                name=name,
                addr=int(sec["sh_addr"]),
                offset=int(sec["sh_offset"]),
                size=int(sec["sh_size"]),
                flags=_section_flags(sec),
                entropy=round(_entropy(data), 4),
            )
        )

    return Report(
        path=str(path),
        size=len(view.data),
        sha256=hashlib.sha256(view.data).hexdigest(),
        elf_machine=str(view.elf["e_machine"]),
        entry=int(view.elf.header["e_entry"]),
        sections=sections,
        defined_function_count=len(defined_funcs),
        imported_function_count=len(imported_funcs),
        exports_interesting=[s for s in defined_funcs if interesting_re.search(s.name)][:120],
        imports_interesting=[s for s in imported_funcs if interesting_re.search(s.name)][:160],
        string_hits=_string_hits(view, text_insns),
        pointer_runs=_pointer_runs(view),
        plt_xrefs_interesting=_plt_xrefs(view, text_insns),
        jni_onload=jni,
        jni_onload_disasm=jni_disasm,
        ollvm_metrics=_ollvm_metrics(text_insns),
    )


def _hx(value: int) -> str:
    return f"0x{value:x}"


def write_markdown(report: Report, path: Path) -> None:
    lines: list[str] = []
    lines.append("# libweapon.so OLLVM 静态初筛报告")
    lines.append("")
    lines.append(f"- 文件：`{Path(report.path).name}`")
    lines.append(f"- 大小：{report.size} bytes")
    lines.append(f"- SHA256：`{report.sha256}`")
    lines.append(f"- 架构：{report.elf_machine}")
    lines.append(f"- Entry：`{_hx(report.entry)}`")
    lines.append(f"- 定义函数数：{report.defined_function_count}")
    lines.append(f"- 导入函数数：{report.imported_function_count}")
    lines.append("")
    lines.append("## OLLVM / kcode-guard 指标")
    m = report.ollvm_metrics
    lines.extend(
        [
            f"- .text 指令数：{m.text_instruction_count}",
            f"- 分支数：{m.branch_count}",
            f"- 条件分支数：{m.cond_branch_count}",
            f"- `br` 间接跳转数：{m.indirect_branch_count}",
            f"- `blr` 间接调用数：{m.indirect_call_count}",
            f"- `movk` 常量拼装数：{m.movk_count}",
            f"- `adrp` 页寻址数：{m.adrp_count}",
            f"- 可疑 dispatch window 数：{len(m.suspicious_dispatch_windows)}",
        ]
    )
    lines.append("")
    lines.append("## 指针表候选")
    if report.pointer_runs:
        for run in report.pointer_runs[:20]:
            target_preview = ", ".join(_hx(v) for v in run.targets[:8])
            lines.append(
                f"- {run.section} `{_hx(run.vaddr)}` file `{_hx(run.file_offset)}` "
                f"count={run.count} targets={target_preview}"
            )
    else:
        lines.append("- 未发现连续 .text 指针表候选")
    lines.append("")
    lines.append("## JNI_OnLoad 反汇编窗口")
    if report.jni_onload:
        lines.append(f"- `JNI_OnLoad`：`{_hx(report.jni_onload.value)}` size={report.jni_onload.size}")
        lines.append("```asm")
        for item in report.jni_onload_disasm[:180]:
            lines.append(f"{_hx(item.address)}: {item.mnemonic} {item.op_str}".rstrip())
        lines.append("```")
    else:
        lines.append("- 未找到 JNI_OnLoad 导出")
    lines.append("")
    lines.append("## 关键字符串命中与直连引用")
    if report.string_hits:
        for hit in report.string_hits[:80]:
            refs = ", ".join(_hx(v) for v in hit.direct_adrp_refs[:8]) or "<none>"
            lines.append(
                f"- `{hit.pattern}` vaddr=`{_hx(hit.vaddr)}` off=`{_hx(hit.offset)}` "
                f"adrp_refs={refs} text=`{hit.text}`"
            )
    else:
        lines.append("- 未命中关键明文字符串")
    lines.append("")
    lines.append("## PLT 调用交叉引用")
    if report.plt_xrefs_interesting:
        for row in report.plt_xrefs_interesting:
            call_preview = ", ".join(_hx(v) for v in row.calls[:12]) or "<none>"
            lines.append(f"- `{row.name}` plt=`{_hx(row.plt)}` calls={len(row.calls)} at {call_preview}")
    else:
        lines.append("- 未识别到直接 PLT 调用")
    lines.append("")
    lines.append("## 关键导入")
    if report.imports_interesting:
        for sym in report.imports_interesting:
            lines.append(f"- `{sym.name}`")
    else:
        lines.append("- 无关键导入命中")
    lines.append("")
    lines.append("## 高熵/可疑节区")
    for sec in sorted(report.sections, key=lambda s: s.entropy, reverse=True)[:16]:
        lines.append(
            f"- `{sec.name}` addr=`{_hx(sec.addr)}` off=`{_hx(sec.offset)}` "
            f"size=`{_hx(sec.size)}` flags={sec.flags} entropy={sec.entropy}"
        )
    lines.append("")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description="libweapon.so OLLVM/kcode-guard 静态初筛")
    parser.add_argument("so", type=Path)
    parser.add_argument("--json-out", type=Path, default=Path("out/libweapon_ollvm_report.json"))
    parser.add_argument("--md-out", type=Path, default=Path("out/libweapon_ollvm_report.md"))
    args = parser.parse_args()

    report = analyze(args.so)
    args.json_out.parent.mkdir(parents=True, exist_ok=True)
    args.json_out.write_text(json.dumps(asdict(report), ensure_ascii=False, indent=2), encoding="utf-8")
    write_markdown(report, args.md_out)

    print(f"path={Path(report.path).name} size={report.size} sha256={report.sha256[:16]}")
    print(f"defined_funcs={report.defined_function_count} imports={report.imported_function_count}")
    print(
        "ollvm "
        f"insns={report.ollvm_metrics.text_instruction_count} "
        f"br={report.ollvm_metrics.indirect_branch_count} "
        f"blr={report.ollvm_metrics.indirect_call_count} "
        f"movk={report.ollvm_metrics.movk_count} "
        f"dispatch_windows={len(report.ollvm_metrics.suspicious_dispatch_windows)}"
    )
    if report.pointer_runs:
        top = report.pointer_runs[0]
        print(f"top_ptr_run={top.section}@0x{top.vaddr:x} count={top.count} first=0x{top.targets[0]:x}")
    print(f"json={args.json_out}")
    print(f"md={args.md_out}")
    return 0


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