# -*- coding: utf-8 -*-
"""Dump JNI/native methods for sdkutils crypto + register-related classes."""
import re
import subprocess
import zipfile
from pathlib import Path
from typing import List

APK = Path(__file__).resolve().parent.parent / "apks" / "seabank_ph_base.apk"
DEXDUMP = Path(r"C:\Users\Administrator\AppData\Local\Android\Sdk\build-tools\37.0.0\dexdump.exe")
OUT = Path(__file__).resolve().parent.parent / "output" / "crypto_jni_targets.txt"

CLASS_PATTERNS = [
    rb"Lcom/shopee/bke/lib/jni/[^;]{1,80};",
    rb"Lcom/shopee/shpssdkbank/wvvvuwwu;",
    rb"Lcom/shopee/shpssdkbank/uwuvuvvww/vvuuuuvvv;",
    rb"Lcom/shopee/shpssdkbank/uwuvuvvww/[^;]{1,40};",
]

EXTRA_KEYWORDS = (
    b"NativeEncrypt",
    b"CharacterCrypto",
    b"SecurityMain",
    b"encrypt",
    b"decrypt",
    b"register",
)


def dump_class(out: str, target: str) -> List[str]:
    lines: list[str] = []
    cap = False
    for line in out.splitlines():
        if f"Class descriptor  : '{target}'" in line:
            cap = True
        elif cap and line.startswith("  Class descriptor") and target not in line:
            break
        if cap:
            lines.append(line)
    return lines


def main() -> None:
    with zipfile.ZipFile(str(APK)) as zf:
        data = b"".join(zf.read(n) for n in zf.namelist() if n.endswith(".dex"))

    classes: set[str] = set()
    for pat in CLASS_PATTERNS:
        for m in re.finditer(pat, data):
            classes.add(m.group().decode()[1:-1].replace("/", "."))

    for kw in EXTRA_KEYWORDS:
        for m in re.finditer(rb"Lcom/[^;]{0,120};", data):
            s = m.group()
            if kw in s or kw in data[data.find(s) : data.find(s) + 4000]:
                classes.add(s.decode()[1:-1].replace("/", "."))

    # log tags -> classes from crash log
    for tag in [
        "com.shopee.bke.lib.jni.utils.f",  # SoUtils
        "com.shopee.bke.lib.jni.uwuwuwuw",
        "com.shopee.bke.lib.jni.uvuvuvuv",
        "com.shopee.bke.lib.jni.uvwwwwuv",
    ]:
        classes.add(tag)

    native_entries: List[str] = []
    with zipfile.ZipFile(str(APK)) as zf:
        for dex_name in zf.namelist():
            if not dex_name.endswith(".dex"):
                continue
            dex = zf.read(dex_name)
            hits = [c for c in classes if c.replace(".", "/") in dex.decode("latin1", errors="ignore")]
            if not hits:
                continue
            tmp = Path(__file__).resolve().parent.parent / "tmp" / "tmp_crypto_jni.dex"
            tmp.write_bytes(dex)
            out = subprocess.check_output(
                [str(DEXDUMP), "-d", str(tmp)], universal_newlines=True, errors="replace"
            )
            for cls in sorted(hits):
                block = dump_class(out, "L" + cls.replace(".", "/") + ";")
                if not block:
                    continue
                native_entries.append(f"\n=== {dex_name} {cls} ===")
                for line in block:
                    if any(
                        k in line
                        for k in ("NATIVE", "name          :", "type          :", "loadLibrary")
                    ):
                        native_entries.append(line)

    text = "\n".join(native_entries)
    OUT.write_text(text, encoding="utf-8")
    print(text)
    print(f"\nwritten {OUT}")


if __name__ == "__main__":
    main()