tests_unpinner.py

import argparse
import os
import re
import subprocess
import sys
from pathlib import Path


def repo_root() -> Path:
    return Path(__file__).resolve().parent


def is_windows() -> bool:
    return sys.platform == "win32"


def venv_bin() -> Path:
    return repo_root() / (".venv/Scripts" if is_windows() else ".venv/bin")


def cli_cmd() -> list[str]:
    exe = venv_bin() / ("android-unpinner.exe" if is_windows() else "android-unpinner")
    if exe.exists():
        return [str(exe)]
    # Fallback to module invocation
    py = venv_bin() / ("python.exe" if is_windows() else "python")
    return [str(py), "-m", "android_unpinner"]


def adb_path() -> Path:
    base = repo_root() / "android_unpinner" / "vendor" / "platform_tools"
    if is_windows():
        return base / "win32" / "adb.exe"
    elif sys.platform == "darwin":
        return base / "darwin" / "adb"
    else:
        return base / "linux" / "adb"


def apksigner_cmd() -> list[str] | None:
    base = repo_root() / "android_unpinner" / "vendor" / "build_tools"
    if is_windows():
        bat = base / "win32" / "apksigner.bat"
        return [str(bat)] if bat.exists() else None
    elif sys.platform == "darwin":
        tool = base / "darwin" / "apksigner"
        return [str(tool)] if tool.exists() else None
    else:
        tool = base / "linux" / "apksigner"
        return [str(tool)] if tool.exists() else None


def run(cmd: list[str], timeout: int | None = None, env: dict | None = None) -> subprocess.CompletedProcess[str]:
    return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env)


def check(proc: subprocess.CompletedProcess[str], desc: str) -> None:
    if proc.returncode != 0:
        raise SystemExit(f"FAIL: {desc}\nRC={proc.returncode}\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}")


def detect_device(serial: str | None) -> str:
    if serial:
        return serial
    out = run([str(adb_path()), "devices", "-l"])  # no check; parse output
    if out.returncode != 0:
        raise SystemExit(f"adb devices failed:\n{out.stderr}")
    lines = [l.strip() for l in out.stdout.splitlines() if l.strip()]
    entries = [l for l in lines[1:] if re.search(r"\sdevice(\s|$)", l)]
    if len(entries) == 1:
        return entries[0].split()[0]
    raise SystemExit(f"Expected exactly one connected device; got: {entries!r}")


def ensure_outdir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--device", "-d", help="Device serial to target")
    p.add_argument("--skip-start", action="store_true", help="Skip start-app & all steps that do JDWP")
    p.add_argument("--timeout-start", type=int, default=25, help="Timeout in seconds for start-app/all")
    args = p.parse_args()

    repo = repo_root()
    cli = cli_cmd()
    adb = adb_path()
    signer = apksigner_cmd()
    apk_in = repo / "httptoolkit-pinning-demo.apk"
    apk_out = repo / "httptoolkit-pinning-demo.unpinned.apk"
    pkg = "tech.httptoolkit.pinning_demo"
    outdir = repo / "tmp_apks_test" / "Chrome"

    print("== Check adb version & devices ==")
    r = run([str(adb), "version"])
    check(r, "adb version")
    print(r.stdout.strip())

    r = run([str(adb), "devices", "-l"])
    check(r, "adb devices -l")
    print(r.stdout.strip())

    serial = detect_device(args.device)
    print(f"Using device: {serial}")

    print("== package-name ==")
    r = run([*cli, "package-name", str(apk_in)])
    check(r, "package-name")
    pn = r.stdout.strip()
    print(pn)
    if pn != pkg:
        raise SystemExit(f"Unexpected package-name: {pn}")

    print("== patch-apks (fresh) ==")

    def verify_apk(path: Path) -> bool | None:
        if signer is None:
            return None
        res = run([*signer, "verify", "--verbose", "--print-certs", str(path)])
        return res.returncode == 0

    existing_valid = apk_out.exists() and (verify_apk(apk_out) is True)
    if existing_valid:
        print(f"Using existing valid patched APK: {apk_out.name}")
    else:
        if apk_out.exists():
            try:
                apk_out.unlink()
            except Exception as e:
                # File locked; if invalid, we must stop. If we couldn't verify, warn.
                ver = verify_apk(apk_out)
                if ver is False:
                    raise SystemExit(
                        f"Existing {apk_out.name} appears invalid and is locked.\n"
                        "Close any apps using it (Explorer preview, antivirus, adb) and rerun.\n"
                        f"Details: {e}"
                    )
                else:
                    print(
                        f"Warning: could not remove existing {apk_out.name} ({e}).\n"
                        "Proceeding with existing file."
                    )
        if not apk_out.exists():
            r = run([*cli, "patch-apks", str(apk_in)])
            check(r, "patch-apks")
            if not apk_out.exists():
                raise SystemExit("Patched APK was not created")
        print("Patched:", apk_out.name)

    if signer is not None:
        print("== verify patched signing (apksigner) ==")
        r = run([*signer, "verify", "--verbose", "--print-certs", str(apk_out)])
        check(r, "apksigner verify")
        print(r.stdout.strip().splitlines()[0])
    else:
        print("apksigner not found; skipping signing verification")

    print("== install patched apk ==" )
    r = run([*cli, "install", "-f", "-d", serial, str(apk_out)])
    check(r, "install")
    print(r.stdout.strip())

    print("== push-resources ==")
    r = run([*cli, "push-resources", "-d", serial])
    check(r, "push-resources")
    print(r.stdout.strip())

    print("== list-packages (check presence) ==")
    r = run([*cli, "list-packages", "-d", serial])
    check(r, "list-packages")
    if pkg not in r.stdout.split():
        raise SystemExit(f"Package {pkg} not found in device list")
    print("Found:", pkg)

    print("== get-apks (com.android.chrome) ==")
    ensure_outdir(outdir)
    r = run([*cli, "get-apks", "-f", "-d", serial, "com.android.chrome", str(outdir)])
    check(r, "get-apks com.android.chrome")
    pulled = outdir / "Chrome.apk"
    if not pulled.exists() or pulled.stat().st_size == 0:
        raise SystemExit("Failed to pull Chrome.apk from device")
    print("Pulled:", pulled)

    if not args.skip_start:
        print("== start-app (with timeout) ==")
        r = run([*cli, "start-app", "-vv", "-d", serial, pkg], timeout=max(5, args.timeout_start))
        # start-app may stall due to JDWP handshake; treat timeout as acceptable
        if r.returncode != 0:
            print("start-app returned non-zero (likely handshake/timeout). Output below:")
            print(r.stdout)
            print(r.stderr)
        else:
            print(r.stdout.strip())

        print("== all (with timeout) ==")
        r = run([*cli, "all", "-vv", "-f", "-d", serial, str(apk_in)], timeout=max(5, args.timeout_start))
        if r.returncode != 0:
            print("all returned non-zero (likely handshake/timeout). Output below:")
            print(r.stdout)
            print(r.stderr)
        else:
            print(r.stdout.strip())

    print("SUCCESS: validation steps completed")


if __name__ == "__main__":
    main()