#!/usr/bin/env bash
# Fast APK protocol extraction via baksmali — bypasses jadx slowness on big/obfuscated APKs.
# Usage: ./decompile-baksmali.sh <package> [out_dir]
#
# Pulls the APK from an ADB-connected device, disassembles with baksmali,
# and extracts every fill-array-data block (the literal int[] byte arrays
# that the app writes to BLE characteristics). Each array is mapped to
# its enclosing function name.
#
# Faster than decompile-app.sh (jadx) for APKs > 30MB because baksmali
# doesn't re-decompile to Java — it works directly on the DEX.

set -e

PKG="${1:?usage: $0 <package> [out_dir]}"
OUT="${2:-/tmp/${PKG}-baksmali}"
APK="/tmp/${PKG}.apk"

ADB_OPTS=""
if [ -n "$ADB_TARGET" ]; then
  ADB_OPTS="-s $ADB_TARGET"
fi

# 1. Find APK on device
echo "==> Locating $PKG on device..."
APK_PATH=$(adb $ADB_OPTS shell pm path "$PKG" | sed 's/package://' | tr -d '\r' | head -1)
if [ -z "$APK_PATH" ]; then
  echo "ERROR: $PKG not installed on device"
  exit 1
fi
echo "    Found: $APK_PATH"

# 2. Pull APK
echo "==> Pulling APK..."
adb $ADB_OPTS pull "$APK_PATH" "$APK"
echo "    Saved to: $APK"

# 3. Get baksmali if missing
if [ ! -f /tmp/baksmali.jar ]; then
  echo "==> Downloading baksmali 2.5.2..."
  curl -sL "https://bitbucket.org/JesusFreke/smali/downloads/baksmali-2.5.2.jar" -o /tmp/baksmali.jar
fi

# 4. Unzip dex
echo "==> Unzipping DEX..."
mkdir -p "$OUT"
unzip -o "$APK" "classes*.dex" -d "$OUT" > /dev/null

# 5. Disassemble
echo "==> Disassembling..."
for dex in "$OUT"/classes*.dex; do
  base=$(basename "$dex" .dex)
  java -jar /tmp/baksmali.jar d "$dex" -o "$OUT/$base" 2>&1 | tail -2
done

# 6. Find the protocol class
echo "==> Searching for protocol classes..."
PROTO_CLASS=$(find "$OUT" -name "NetConnectBle.smali" -o -name "*Ble*Service*.smali" -o -name "*Command*.smali" 2>/dev/null | head -3)
if [ -z "$PROTO_CLASS" ]; then
  echo "    No obvious protocol class. Try manual grep:"
  echo "    grep -rl 'fill-array-data' $OUT/ | head"
  exit 0
fi
echo "    Found: $PROTO_CLASS"
echo

# 7. Extract every fill-array-data with surrounding context
echo "==> Protocol bytes extracted from smali:"
for f in $PROTO_CLASS; do
  echo
  echo "    ==================== $f ===================="
  python3 << PYEOF
import re
content = open("$f").read()

# Split on .method declarations to get function → body mapping
methods = re.split(r'\n\.method ', content)
for chunk in methods[1:]:
    # First line is the method signature
    sig = chunk.split('\n', 1)[0]
    body = chunk.split('\n', 1)[1] if '\n' in chunk else ''
    # Find all array-data blocks within this method body
    arrays = re.findall(
        r'fill-array-data\s+\w+,\s+:(\w+).*?\.array-data\s+\d+\n((?:\s+0x[0-9a-f]+\n)+)\.end array-data',
        body, re.DOTALL
    )
    for name, body_arr in arrays:
        bytes_list = [int(b.strip(), 16) for b in body_arr.strip().split('\n') if b.strip()]
        bytes_str = ' '.join(f'{b:02X}' for b in bytes_list)
        print(f'  {sig.strip()[:60]}')
        print(f'    array {name}: {bytes_str}')
        print()
PYEOF
done