"""
extract_apk_protocol.py - Extract BLE UUIDs and protocol from KAIYU APK

Pulls the APK directly from the connected Android device via ADB, then
searches the DEX bytecode for UUID patterns and BLE write calls.
No jadx or external decompiler needed — pure Python + adb.

Usage:
    python sniffer/extract_apk_protocol.py
"""

import subprocess
import zipfile
import re
import os
import sys
import json

ADB       = r"D:\platform-tools\adb.exe"
PKG       = "com.qunyu.kaiyu"
OUT_DIR   = r"D:\Claude\bugreport\apk_extract"
APK_LOCAL = os.path.join(OUT_DIR, "kaiyu.apk")
RESULTS   = os.path.join(OUT_DIR, "protocol_hints.json")

UUID_RE = re.compile(
    r'[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}'
)

KNOWN_UUIDS = {
    "0000ffe1-0000-1000-8000-00805f9b34fb": "LEDBLE/MagicLight write",
    "0000fff3-0000-1000-8000-00805f9b34fb": "ELK-BLEDOM write",
    "0000ffd9-0000-1000-8000-00805f9b34fb": "LED controller write",
    "0000ffe9-0000-1000-8000-00805f9b34fb": "MagicLight older write",
    "0000ffb2-0000-1000-8000-00805f9b34fb": "Colorific write",
    "0000ffe0-0000-1000-8000-00805f9b34fb": "LEDBLE service",
    "0000fff0-0000-1000-8000-00805f9b34fb": "ELK-BLEDOM service",
}

BLE_KEYWORDS = [
    "writecharacteristic", "writegattchar", "blewrite", "writevalue",
    "setcolor", "set_color", "setrgb", "rgb", "brightness",
    "ffe1", "fff3", "ffd9", "ffe0", "fff0", "ffe9",
    "7e00", "cc23", "cc24", "56ff",
    "service_uuid", "characteristic_uuid", "gatt",
    "BluetoothGattCharacteristic", "BluetoothLeService",
]


def run(cmd):
    result = subprocess.run(cmd, capture_output=True, text=True)
    return result.stdout.strip(), result.returncode


def pull_apk() -> bool:
    os.makedirs(OUT_DIR, exist_ok=True)

    if os.path.exists(APK_LOCAL) and os.path.getsize(APK_LOCAL) > 100_000:
        print(f"APK already cached: {APK_LOCAL}  ({os.path.getsize(APK_LOCAL):,} bytes)")
        return True

    print(f"Finding APK on device for package: {PKG}")
    out, rc = run([ADB, "shell", f"pm path {PKG}"])
    if rc != 0 or "package:" not in out:
        print(f"[ERROR] Package not found on device: {out}")
        print("Make sure the KAIYU app is installed and the phone is connected via USB.")
        return False

    apk_paths = [line.replace("package:", "").strip() for line in out.splitlines()]
    apk_path  = next((p for p in apk_paths if "base" in p.lower()), apk_paths[0])
    print(f"APK on device: {apk_path}")

    print(f"Pulling APK... (this may take 10-20 seconds)")
    _, rc = run([ADB, "pull", apk_path, APK_LOCAL])
    if rc != 0 or not os.path.exists(APK_LOCAL):
        # Try copy to sdcard first (some devices block direct pull)
        print("Direct pull failed, trying via sdcard...")
        run([ADB, "shell", f"cp {apk_path} /sdcard/kaiyu_tmp.apk"])
        _, rc = run([ADB, "pull", "/sdcard/kaiyu_tmp.apk", APK_LOCAL])
        run([ADB, "shell", "rm /sdcard/kaiyu_tmp.apk"])

    if not os.path.exists(APK_LOCAL) or os.path.getsize(APK_LOCAL) < 1000:
        print("[ERROR] APK pull failed.")
        return False

    print(f"APK pulled: {os.path.getsize(APK_LOCAL):,} bytes")
    return True


def extract_strings_from_dex(apk_path: str) -> list[str]:
    """Extract all printable ASCII strings >= 6 chars from DEX bytecode."""
    strings = []
    with zipfile.ZipFile(apk_path, 'r') as zf:
        dex_files = [n for n in zf.namelist() if n.endswith('.dex')]
        print(f"DEX files in APK: {dex_files}")
        for dex_name in dex_files:
            raw = zf.read(dex_name)
            current = []
            for byte in raw:
                if 32 <= byte < 127:
                    current.append(chr(byte))
                else:
                    if len(current) >= 6:
                        strings.append("".join(current))
                    current = []
            if len(current) >= 6:
                strings.append("".join(current))

    print(f"Total strings extracted: {len(strings):,}")
    return strings


def analyse(strings: list[str]) -> dict:
    results = {
        "full_uuids": {},
        "known_uuid_matches": [],
        "ble_strings": [],
        "resource_uuids": [],
    }

    for s in strings:
        for uuid in UUID_RE.findall(s):
            ul = uuid.lower()
            results["full_uuids"][ul] = results["full_uuids"].get(ul, 0) + 1
            if ul in KNOWN_UUIDS:
                hit = f"{ul}  ->  {KNOWN_UUIDS[ul]}"
                if hit not in results["known_uuid_matches"]:
                    results["known_uuid_matches"].append(hit)

        slower = s.lower()
        if any(kw in slower for kw in BLE_KEYWORDS) and len(s) < 300:
            if s not in results["ble_strings"]:
                results["ble_strings"].append(s)

    return results


def check_resources(apk_path: str) -> list[str]:
    found = []
    with zipfile.ZipFile(apk_path, 'r') as zf:
        for name in zf.namelist():
            if name.startswith("res/") and name.endswith(".xml"):
                try:
                    content = zf.read(name).decode("utf-8", errors="ignore")
                    for uuid in UUID_RE.findall(content):
                        found.append(f"{name}: {uuid}")
                except Exception:
                    pass
    return found


def main():
    print("=" * 60)
    print("  KAIYU APK Protocol Extractor")
    print("=" * 60)
    print()

    if not pull_apk():
        sys.exit(1)

    print("\nExtracting strings from DEX bytecode...")
    strings = extract_strings_from_dex(APK_LOCAL)

    print("Analysing for BLE UUIDs and protocol hints...")
    results = analyse(strings)

    print("\nChecking XML resources...")
    results["resource_uuids"] = check_resources(APK_LOCAL)

    # ── Print results ──────────────────────────────────────────────────────
    print(f"\n{'='*60}")
    print(f"  RESULTS")
    print(f"{'='*60}")

    print(f"\n--- Known BLE LED controller UUIDs ({len(results['known_uuid_matches'])}) ---")
    if results["known_uuid_matches"]:
        for m in results["known_uuid_matches"]:
            print(f"  *** {m}")
    else:
        print("  None of the well-known LED UUIDs found as plain text in DEX")

    print(f"\n--- All UUIDs found ({len(results['full_uuids'])}) ---")
    for uuid, count in sorted(results["full_uuids"].items(), key=lambda x: -x[1]):
        known = f"  <-- {KNOWN_UUIDS[uuid]}" if uuid in KNOWN_UUIDS else ""
        print(f"  {uuid}  (x{count}){known}")

    print(f"\n--- BLE-related strings ({len(results['ble_strings'])}) ---")
    for s in results["ble_strings"][:60]:
        print(f"  {repr(s)}")

    if results["resource_uuids"]:
        print(f"\n--- Resource file UUIDs ({len(results['resource_uuids'])}) ---")
        for r in results["resource_uuids"]:
            print(f"  {r}")

    # Save full results
    with open(RESULTS, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nFull results saved: {RESULTS}")

    # ── Recommendation ─────────────────────────────────────────────────────
    print(f"\n{'='*60}")
    print(f"  NEXT STEP")
    print(f"{'='*60}")
    if results["known_uuid_matches"]:
        print("  Known UUID found! Update controller/config.py with the write UUID.")
        print("  Then run: python scanner/ble_scanner.py --probe <device_address>")
    elif results["full_uuids"]:
        print("  Custom UUIDs found. Try each as WRITE_CHARACTERISTIC in config.py")
        print("  and test: python scanner/ble_scanner.py --probe <device_address>")
    else:
        print("  No UUIDs in DEX (may be obfuscated or in native lib).")
        print("  Run the interactive probe instead:")
        print("  python scanner/ble_scanner.py --probe <device_address>")


if __name__ == "__main__":
    main()