#!/usr/bin/env python3
"""dex 静态判定: 扫描 APK 全部 classes*.dex, 定位目标类及其方法的 access_flags。

用途: 确定 NativeBridge.b / WupHelper.getGuid 等是否声明为 native (ACC_NATIVE=0x100),
完全不碰设备。
"""
import struct
import sys
from pathlib import Path

ACC_PUBLIC = 0x1
ACC_PRIVATE = 0x2
ACC_PROTECTED = 0x4
ACC_STATIC = 0x8
ACC_FINAL = 0x10
ACC_SYNCHRONIZED = 0x20
ACC_VOLATILE = 0x40
ACC_NATIVE = 0x100
ACC_ABSTRACT = 0x400
ACC_CONSTRUCTOR = 0x10000
ACC_DECLARED_SYNCHRONIZED = 0x20000

ACC_NAMES = {
    ACC_PUBLIC: "public", ACC_PRIVATE: "private", ACC_PROTECTED: "protected",
    ACC_STATIC: "static", ACC_FINAL: "final", ACC_SYNCHRONIZED: "synchronized",
    ACC_VOLATILE: "volatile", ACC_NATIVE: "native", ACC_ABSTRACT: "abstract",
    ACC_CONSTRUCTOR: "constructor", ACC_DECLARED_SYNCHRONIZED: "declared-synchronized",
}


def read_uleb128(buf, off):
    result = 0
    shift = 0
    while True:
        b = buf[off]
        off += 1
        result |= (b & 0x7F) << shift
        if (b & 0x80) == 0:
            break
        shift += 7
    return result, off


def read_sleb128(buf, off):
    result = 0
    shift = 0
    b = 0
    while True:
        b = buf[off]
        off += 1
        result |= (b & 0x7F) << shift
        shift += 7
        if (b & 0x80) == 0:
            break
    if b & 0x40:
        result |= -(1 << shift)
    return result, off


class Dex:
    def __init__(self, data):
        self.data = data
        assert data[:8] == b"dex\n035\0" or data[:8] == b"dex\n037\0" or data[:8] == b"dex\n038\0" or data[:8] == b"dex\n039\0", data[:8]
        self.string_ids_size, self.string_ids_off = struct.unpack_from("<II", data, 0x38)
        self.type_ids_size, self.type_ids_off = struct.unpack_from("<II", data, 0x40)
        self.proto_ids_size, self.proto_ids_off = struct.unpack_from("<II", data, 0x48)
        self.field_ids_size, self.field_ids_off = struct.unpack_from("<II", data, 0x50)
        self.method_ids_size, self.method_ids_off = struct.unpack_from("<II", data, 0x58)
        self.class_defs_size, self.class_defs_off = struct.unpack_from("<II", data, 0x60)

    def string(self, idx):
        off = struct.unpack_from("<I", self.data, self.string_ids_off + idx * 4)[0]
        # skip uleb128 utf16 size
        _, off = read_uleb128(self.data, off)
        end = self.data.index(b"\0", off)
        return self.data[off:end].decode("utf-8", "replace")

    def type(self, idx):
        desc_idx = struct.unpack_from("<I", self.data, self.type_ids_off + idx * 4)[0]
        return self.string(desc_idx)

    def method(self, idx):
        class_idx, proto_idx, name_idx = struct.unpack_from("<HHI", self.data, self.method_ids_off + idx * 8)
        return self.type(class_idx), self.string(name_idx), proto_idx

    def proto(self, idx):
        shorty_idx, return_type_idx, params_off = struct.unpack_from("<III", self.data, self.proto_ids_off + idx * 12)
        rt = self.type(return_type_idx)
        params = []
        if params_off:
            size = struct.unpack_from("<I", self.data, params_off)[0]
            for i in range(size):
                params.append(self.type(struct.unpack_from("<H", self.data, params_off + 4 + i * 2)[0]))
        return rt, params

    def class_defs_iter(self):
        for i in range(self.class_defs_size):
            fields = struct.unpack_from("<IIIIIIII", self.data, self.class_defs_off + i * 32)
            class_idx, access_flags, superclass_idx, interfaces_off, source_file_idx, annotations_off, class_data_off, static_values_off = fields
            yield {
                "class_idx": class_idx,
                "access_flags": access_flags,
                "class_data_off": class_data_off,
            }

    def class_data_methods(self, cd_off):
        """return list of (name_idx, access_flags, code_off)"""
        if cd_off == 0:
            return []
        off = cd_off
        static_fields, off = read_uleb128(self.data, off)
        instance_fields, off = read_uleb128(self.data, off)
        direct_methods, off = read_uleb128(self.data, off)
        virtual_methods, off = read_uleb128(self.data, off)
        # skip field entries (each = field_idx_diff + access_flags)
        fidx = 0
        for _ in range(static_fields + instance_fields):
            diff, off = read_uleb128(self.data, off)
            fidx += diff
            _, off = read_uleb128(self.data, off)
        res = []
        idx = 0
        for _ in range(direct_methods + virtual_methods):
            idx_diff, off = read_uleb128(self.data, off)
            idx += idx_diff
            acc, off = read_uleb128(self.data, off)
            code_off, off = read_uleb128(self.data, off)
            res.append((idx, acc, code_off))
        return res

    def flags_str(self, flags):
        parts = []
        for bit, name in ACC_NAMES.items():
            if flags & bit:
                parts.append(name)
        return " ".join(parts) if parts else ""


def main():
    dex_dir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("work/apk_extract")
    # 目标: 类名 -> 关注的方法名集合
    targets = {
        "Lcom/huya/security/hydeviceid/NativeBridge;": {"b", "klog", "a", "c", "h", "g"},
        "Lcom/huya/security/hydeviceid/NativeEntry;": {"getGUID", "init", "getMID", "getHDID", "getCDID", "getSDID"},
        "Lcom/huya/security/hydeviceid/WupHelper;": {"getGuid", "getGuid1", "buildGuid"},
        "Lcom/huya/security/hydeviceid/HuyaAuthCore;": None,  # 全部方法
    }
    found = {}
    for dex_path in sorted(dex_dir.glob("classes*.dex")):
        data = dex_path.read_bytes()
        dex = Dex(data)
        for cd in dex.class_defs_iter():
            try:
                cls = dex.type(cd["class_idx"])
            except struct.error:
                continue
            if cls not in targets:
                continue
            wanted = targets[cls]
            try:
                methods = dex.class_data_methods(cd["class_data_off"])
            except struct.error as e:
                print(f"  [skip] {cls} in {dex_path.name}: class_data parse error {e}")
                continue
            entry = found.setdefault(cls, [])
            entry.append(str(dex_path.name))
            out = []
            for m_idx, acc, code_off in methods:
                if m_idx >= dex.method_ids_size:
                    continue
                mcls, mname, proto_idx = dex.method(m_idx)
                if wanted is not None and mname not in wanted and not (wanted is None):
                    continue
                rt, params = dex.proto(proto_idx)
                sig = f"({''.join(params)}){rt}"
                out.append(f"  {dex.flags_str(acc):32s} {mname}{sig}  code_off={code_off:#x}")
            if out:
                print(f"== {cls}  <in {dex_path.name}>  class_flags={dex.flags_str(cd['access_flags'])}")
                print("\n".join(out))
    if not found:
        # 打印包含这些类名串的 dex, 方便定位动态下载的 dex
        print("target classes NOT found in any classes*.dex; searching string blobs...")
        for cls in targets:
            blob = cls.encode()
            for dex_path in sorted(dex_dir.glob("classes*.dex")):
                data = dex_path.read_bytes()
                if blob in data:
                    print(f"  {cls} string present in {dex_path.name}")


if __name__ == "__main__":
    main()