main.py

#!/usr/bin/env python3
"""
TrashDroid — Automated Android DAST Framework
Main entry point and phase orchestrator.

Usage:
    python main.py
    python main.py --skip-preflight
    python main.py --phases 1,3,5,8
"""

from __future__ import annotations

import argparse
import atexit
import os
import signal
import sys
import threading
import traceback
from pathlib import Path

from rich.align import Align
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Confirm

from core import __version__
from core.adb import ADB
from core.config import BANNER, TIMING, Config
from core.drozer import Drozer
from core.pii_runtime import initialize_pii_detection
from core.report import ReportGenerator
from core.runtime_cleanup import RuntimeCleanupManager
from core.screenshot import ScreenshotManager
from phases.backup import run_backup_analysis
from phases.drozer_testing import run_drozer_testing
from phases.dump_verify import run_dump_verification
from phases.filesystem import run_filesystem_analysis
from phases.logcat import run_logcat_monitoring
from phases.manifest import run_manifest_analysis
from phases.memory import run_memory_analysis
from phases.post_logout import run_post_logout_testing
from phases.preflight import run_preflight
from phases.runtime_hardening import run_runtime_hardening
from phases.setup import get_apk_input, install_and_prepare, select_device
from utils.helpers import is_valid_package_name

console = Console()

ALL_PHASES = {
    1: ("Phase I   — Drozer Component Testing", "drozer"),
    3: ("Phase III — Local File System Analysis", "filesystem"),
    4: ("Phase IV  — Dump File Verification", "dump_verify"),
    5: ("Phase V   — Logcat Monitoring", "logcat"),
    6: ("Phase VI  — Memory Analysis", "memory"),
    7: ("Phase VII — ADB Backup Analysis", "backup"),
    8: ("Phase VIII— Manifest Analysis", "manifest"),
    9: ("Phase IX  — Post-Logout Access Control", "post_logout"),
    10: ("Phase X   — Runtime Hardening (SSL pinning / root detection)", "runtime_hardening"),
}

_EXAMPLES = """\
examples:
  python main.py                                    interactive (pick device + app)
  python main.py --auto --package com.example.app   non-interactive, all phases
  python main.py --phases 1,3,8 --package com.x      run only the named phases
  python main.py --presidio --auto --package com.x   enable Presidio PII detection
  python main.py --version                           print version and exit
"""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Android DAST — Automated VAPT Framework",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=_EXAMPLES,
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"TrashDroid {__version__}",
    )
    parser.add_argument(
        "--skip-preflight",
        action="store_true",
        help="Skip tool availability checks",
    )
    parser.add_argument(
        "--phases",
        type=str,
        default="",
        help="Comma-separated list of phase numbers to run (e.g. 1,3,5,8). Default: all",
    )
    parser.add_argument(
        "--package",
        type=str,
        default="",
        help="Package name (skip interactive prompt)",
    )
    parser.add_argument(
        "--apk",
        type=str,
        default="",
        help="Path to APK file (skip interactive prompt)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="",
        help="Device serial (skip interactive prompt)",
    )
    parser.add_argument(
        "--auto",
        action="store_true",
        help="Non-interactive mode — uses sensible defaults for all prompts",
    )
    parser.add_argument(
        "--report",
        choices=["client", "internal"],
        default="client",
        help="Report detail level (internal includes prompts)",
    )
    parser.add_argument(
        "--screenshot-delay",
        type=float,
        default=TIMING.screenshot_settle_delay,
        help=f"Delay in seconds before capturing a screenshot (default: {TIMING.screenshot_settle_delay})",
    )
    parser.add_argument(
        "--presidio",
        action="store_true",
        help="Enable Presidio PII detection (regex + checksum validators); falls back to regex-only on init failure",
    )
    parser.add_argument(
        "--ner",
        action="store_true",
        help="Enable GLiNER NER backend for ML-based PII detection (implies --presidio, fails fast on init errors)",
    )
    parser.add_argument(
        "--ai-review",
        action="store_true",
        help="After the run, auto-run `claude` (or $TRASHDROID_REVIEW_CMD) headless over the "
             "ai_review/ package to write final_report.md (+ final_report.html)",
    )
    return parser.parse_args()


def _execute_phase(phase_runner, phase_name: str, phase_num: int, config: Config, args) -> None:
    """Run one phase, applying a wall-clock watchdog in --auto mode.

    Interactive runs call the phase directly (they are user-paced). In --auto mode
    the phase runs in a daemon thread joined with TIMING.phase_budget_sec; if it
    overruns, the thread is abandoned and a "timed out" finding is recorded so a
    single wedged adb/drozer/frida call can't hang the whole run. Exceptions raised
    inside the thread are re-raised on the caller so the existing per-phase
    error handling still applies.
    """
    budget = TIMING.phase_budget_sec
    if not (getattr(args, "auto", False) and budget and budget > 0):
        phase_runner()
        return

    error_holder: dict[str, BaseException] = {}

    def _target() -> None:
        try:
            phase_runner()
        except BaseException as e:  # re-raised on the main thread below
            error_holder["err"] = e

    worker = threading.Thread(target=_target, name=f"phase-{phase_num}", daemon=True)
    worker.start()
    worker.join(timeout=budget)

    if worker.is_alive():
        console.print(
            f"[red]Phase {phase_num} exceeded the {budget:.0f}s auto-mode budget — "
            f"abandoning and continuing.[/red]"
        )
        config.add_finding(
            phase_name,
            "Phase timed out",
            "Info",
            f"Phase {phase_num} exceeded the {budget:.0f}s --auto budget and was abandoned. "
            "A device/tool call likely hung; re-run this phase individually to investigate.",
        )
        return

    if "err" in error_holder:
        raise error_holder["err"]


def main() -> int:
    args = parse_args()

    console.print(Align.left(Panel(BANNER, style="bright_white", expand=True, subtitle="Author: 0xs0m")))

    # ── Pre-flight ──
    if not args.skip_preflight:
        if not run_preflight():
            return 1
    else:
        console.print("[yellow]Pre-flight checks skipped (--skip-preflight).[/yellow]")

    # ── Device selection ──
    if args.device:
        available_devices = ADB.get_devices()
        if args.device not in available_devices:
            console.print(
                f"[red]Device '{args.device}' not found via adb. "
                f"Available: {available_devices or 'none'}[/red]"
            )
            return 1
        device_id = args.device
        console.print(f"[green]Using device: {device_id}[/green]")
    else:
        device_id = select_device()
        if not device_id:
            return 1

    adb = ADB(device_id)
    config = Config(device_id=device_id)

    try:
        device_info = adb.get_device_info()
    except Exception as e:
        console.print(f"[red]Failed to read device information: {e}[/red]")
        return 1
    console.print(
        f"[green]Device: {device_info['model']} | "
        f"Android {device_info['android_version']} | "
        f"SDK {device_info['sdk']}[/green]"
    )

    try:
        root_status = adb.is_rooted()
    except Exception as e:
        console.print(f"[yellow]⚠ Could not determine root status ({e}); assuming not rooted.[/yellow]")
        root_status = False
    if root_status:
        console.print("[green]✓ Device is rooted.[/green]")
    else:
        console.print("[yellow]⚠ Device does not appear to be rooted. Some tests may fail.[/yellow]")

    # Check on-device prerequisites
    from phases.preflight import verify_device_prerequisites
    if not verify_device_prerequisites(adb):
        return 1

    # ── Target app input ──
    if args.package:
        if not is_valid_package_name(args.package):
            console.print(f"[red]Invalid --package '{args.package}'. Expected e.g. com.example.app[/red]")
            return 1
        config.package_name = args.package
        config.apk_path = args.apk or None
        config.is_preinstalled = not bool(args.apk)
        if config.apk_path and not Path(config.apk_path).exists():
            console.print(f"[red]APK file not found: {config.apk_path}[/red]")
            return 1
    elif args.apk:
        # --apk without --package: derive the package id from the APK rather than
        # silently ignoring --apk and re-prompting.
        if not Path(args.apk).exists():
            console.print(f"[red]APK file not found: {args.apk}[/red]")
            return 1
        derived = adb.get_package_name_from_apk(args.apk)
        if not derived or not is_valid_package_name(derived):
            console.print(
                f"[red]Could not determine a valid package name from {args.apk}. "
                f"Pass --package explicitly.[/red]"
            )
            return 1
        config.package_name = derived
        config.apk_path = args.apk
        config.is_preinstalled = False
        console.print(f"[green]Derived package from APK: {derived}[/green]")
    elif args.auto:
        console.print("[red]--auto requires --package or --apk (cannot prompt interactively).[/red]")
        return 1
    else:
        apk_path, pkg, is_pre = get_apk_input(adb)
        config.apk_path = apk_path
        config.package_name = pkg
        config.is_preinstalled = is_pre

    if not config.package_name.strip():
        console.print("[red]Package name cannot be empty.[/red]")
        return 1

    config.auto_mode = args.auto
    config.report_mode = args.report
    config.screenshot_delay = args.screenshot_delay
    config.init_output()

    # ── Initialize PII detection backend (eager warmup) ──
    pii_init_rc = initialize_pii_detection(
        config=config,
        use_presidio=args.presidio,
        use_ner=args.ner,
        console=console,
    )
    if pii_init_rc != 0:
        return pii_init_rc

    # ── Install & prepare ──
    install_and_prepare(adb, config)

    # ── Init helpers ──
    drozer = Drozer(device_id, rooted=root_status)
    screenshotter = ScreenshotManager(adb, config.screenshot_dir, config)

    # ── Register cleanup handler (prevents orphan logcat/scrcpy + saves partial report) ──
    cleanup_manager = RuntimeCleanupManager(
        screenshotter=screenshotter,
        config=config,
        device_info=device_info,
    )

    def _cleanup(generate_partial_report: bool = True):
        cleanup_manager.cleanup(generate_partial_report=generate_partial_report)

    atexit.register(_cleanup)

    def _signal_handler(signum, frame):
        console.print(f"\n[yellow]Received signal {signum} — cleaning up...[/yellow]")
        _cleanup(generate_partial_report=True)
        sys.exit(128 + signum)

    signal.signal(signal.SIGTERM, _signal_handler)
    signal.signal(signal.SIGINT, _signal_handler)

    # Determine which phases to run
    if args.phases:
        selected = set()
        invalid_phases: list[str] = []
        for p in args.phases.split(","):
            try:
                phase_num = int(p.strip())
                if phase_num in ALL_PHASES:
                    selected.add(phase_num)
                else:
                    invalid_phases.append(str(phase_num))
            except ValueError:
                invalid_phases.append(p.strip())
        if invalid_phases:
            console.print(f"[yellow]Ignoring invalid phase(s): {', '.join(invalid_phases)}[/yellow]")
    else:
        selected = set(ALL_PHASES.keys())

    console.print(f"\n[bold]Phases to run:[/bold] {sorted(selected)}\n")

    # ── Optionally start scrcpy for live viewing ──
    if any(p in selected for p in [1, 9]):
        if args.auto:
            console.print("[yellow]Auto-mode: skipping scrcpy live view (screenshots still captured via adb).[/yellow]")
        else:
            console.print(
                "[bold]scrcpy provides a live mirror of the device screen but can be GPU-heavy.[/bold]\n"
                "[dim]Screenshots are always captured in the background via adb regardless of this choice.[/dim]"
            )
            want_scrcpy = Confirm.ask("Launch scrcpy for live screen mirroring?", default=True)
            if want_scrcpy:
                screenshotter.start_scrcpy()

    # ── Start background logcat collector ──
    from utils.logcat_collector import BackgroundLogcatCollector
    bg_logcat = BackgroundLogcatCollector(device_id, config.package_name, config.output_dir)
    bg_logcat.start()
    console.print("[dim]Background logcat collector started.[/dim]")
    cleanup_manager.set_background_collector(bg_logcat)

    # ── Execute phases ──
    phase_runners = {
        1: lambda: run_drozer_testing(config, adb, drozer, screenshotter),
        3: lambda: run_filesystem_analysis(config, adb), #include trufflehog
        4: lambda: run_dump_verification(config, adb),
        5: lambda: run_logcat_monitoring(config, adb),
        6: lambda: run_memory_analysis(config, adb),
        7: lambda: run_backup_analysis(config, adb),
        8: lambda: run_manifest_analysis(config, adb),
        9: lambda: run_post_logout_testing(config, adb, drozer, screenshotter),
        10: lambda: run_runtime_hardening(config, adb, screenshotter),
    }

    for phase_num in sorted(selected):
        if phase_num not in phase_runners:
            console.print(f"[yellow]Unknown phase {phase_num}, skipping.[/yellow]")
            continue

        phase_name = ALL_PHASES[phase_num][0]
        try:
            _execute_phase(phase_runners[phase_num], phase_name, phase_num, config, args)
        except KeyboardInterrupt:
            console.print(f"\n[yellow]Phase {phase_num} interrupted by user.[/yellow]")
            if not args.auto and not Confirm.ask("Continue to next phase?", default=True):
                break
        except Exception as e:
            console.print(f"\n[red]Error in {phase_name}: {e}[/red]")
            console.print(f"[dim]{traceback.format_exc()}[/dim]")
            config.add_finding(
                phase_name,
                "Phase execution error",
                "Info",
                f"Phase {phase_num} encountered an error:\n{traceback.format_exc()}",
            )
            if not args.auto and not Confirm.ask("Continue to next phase?", default=True):
                break

    # ── Stop background logcat collector and integrate findings ──
    bg_logcat.stop()
    bg_findings = bg_logcat.save_and_scan()
    for bf in bg_findings:
        config.add_finding("Background Logcat", bf["title"], bf["severity"], bf["detail"])
    if bg_findings:
        console.print(f"[yellow]Background logcat found {len(bg_findings)} finding(s).[/yellow]")

    # ── Stop scrcpy ──
    screenshotter.stop_scrcpy()

    # ── Generate report ──
    console.print("\n[bold cyan]═══ Generating Report ═══[/bold cyan]\n")
    reporter = ReportGenerator(config, device_info)
    report_path = reporter.generate()
    cleanup_manager.mark_final_report_generated()

    total_findings = sum(len(v) for v in config.findings.values())
    total_screenshots = len(config.screenshots)
    total_commands = len(config.commands_log)

    console.print(Panel(
        f"[bold green]DAST Assessment Complete[/bold green]\n\n"
        f"  Report:      {report_path}\n"
        f"  Findings:    {total_findings}\n"
        f"  Screenshots: {total_screenshots}\n"
        f"  Commands:    {total_commands}\n"
        f"  Output dir:  {config.output_dir}",
        title="Summary",
        style="green",
        expand=False,
    ))

    # ── Assemble the AI-review evidence package (findings + screenshots + raw logs + prompt) ──
    from core.ai_review import (
        assemble_review_package,
        launch_claude_interactive,
        print_next_steps,
        run_claude_review,
    )
    pkg = assemble_review_package(config, device_info, report_path)

    if args.ai_review:                       # explicit flag → headless, unattended
        run_claude_review(pkg, console)
    elif not args.auto:                      # interactive → let the operator choose
        from rich.prompt import Prompt
        console.print(
            "\n[bold]Triage this evidence package with an AI now?[/bold]\n"
            "  [cyan]1[/cyan]) Interactive [white]claude[/white] session  "
            "[dim](recommended — it can ask you to connect the device / log out and verify live)[/dim]\n"
            "  [cyan]2[/cyan]) Headless [white]claude[/white]  "
            "[dim](streaming, unattended — writes final_report.md, no live Q&A)[/dim]\n"
            "  [cyan]3[/cyan]) Custom / cloud command  "
            "[dim]($TRASHDROID_REVIEW_CMD — OpenRouter / Ollama / aider)[/dim]\n"
            "  [cyan]4[/cyan]) Just show me the prompt  [dim](paste into any AI — claude.ai, ChatGPT, …)[/dim]"
        )
        choice = Prompt.ask("Choice", choices=["1", "2", "3", "4"], default="1")
        if choice == "1":
            launch_claude_interactive(pkg, console)
        elif choice == "2":
            run_claude_review(pkg, console)
        elif choice == "3":
            if not os.environ.get("TRASHDROID_REVIEW_CMD"):
                console.print("[yellow]$TRASHDROID_REVIEW_CMD is not set. Set it first, e.g.:[/yellow]\n"
                              "  [white]export TRASHDROID_REVIEW_CMD='aider --message-file {prompt_file} --yes'[/white]\n"
                              "[dim]Then re-run, or use ./run_review.sh in the package.[/dim]")
            else:
                run_claude_review(pkg, console)

    print_next_steps(pkg, config, console)

    return 0


if __name__ == "__main__":
    sys.exit(main())