main.py

#!/usr/bin/env python3
"""
Vulnerability Scanning Orchestrator — Entry Point

Interactive mode (default):
    python main.py

Batch mode (legacy):
    python main.py --batch --targets targets.txt --profile standard
"""

import argparse
import asyncio
import signal
import sys
from pathlib import Path
from urllib.parse import urlparse

import structlog
import subprocess

from orchestrator.config_loader import load_config
from orchestrator.core import Orchestrator
from orchestrator.safety import AllowlistEnforcer

log = structlog.get_logger()

# ── Bonsai Integration ────────────────────────────

def run_bonsai_login() -> None:
    """Run 'bonsai login' and wait for user to authenticate."""
    print("\n  \033[36m[*] Initializing Bonsai Proxy Authentication...\033[0m")
    try:
        # We use subprocess.run with capture_output=False to show the login URL in the terminal
        # The user needs to see the URL and visit it.
        process = subprocess.run(["bonsai", "login"], check=False)
        if process.returncode != 0:
            print("  \033[31m[!] Bonsai login failed. Continuing anyway, but AI may be unavailable.\033[0m")
        else:
            print("  \033[32m[OK] Bonsai authentication initialized.\033[0m")
    except FileNotFoundError:
        print("  \033[33m[!] 'bonsai' CLI not found. AI features might fail if they depend on it.\033[0m")
    except Exception as exc:
        print(f"  \033[31m[!] Error running bonsai login: {exc}\033[0m")
    print()

# ── Pretty banner ─────────────────────────────────
BANNER = r"""
 ___      ___  ___  ___   ___  _  _   ___   ___    _    _  _
\ \ / /  | | | | | |  | \| / __| / __|  /  \ | \| |
 \ V /| |_| | |__| .` \__ \| (__ / __ \| .` |
  \_/  \___/|____|_|\_|___/ \___|_/  \_\_|\_|
         P R O F E S S I O N A L   S C A N N E R
"""


def _print_banner() -> None:
    print("\033[36m" + BANNER + "\033[0m")
    print("  Automated Vulnerability Scanner")
    print("  ─────────────────────────────────────────\n")


def _normalize_target(raw: str) -> str:
    """
    Accept flexible user input and return a clean domain or URL.
    Examples:
        'google.com'            -> 'google.com'
        'https://example.com/'  -> 'example.com'
        'Example App (example.com)' -> 'example.com'
        'http://test.site/path' -> 'test.site'
    """
    raw = raw.strip()
    if not raw:
        return ""

    # If it looks like a URL, extract the hostname
    if "://" in raw:
        parsed = urlparse(raw)
        return parsed.hostname or raw

    # Strip common noise: trailing slashes, www prefix handled later
    raw = raw.strip("/").strip()

    # If it has spaces, try to find a domain-like token
    if " " in raw:
        for token in raw.split():
            token = token.strip("()[]<>,\"'")
            if "." in token and " " not in token:
                return token
        # Fallback: just use the whole thing (will fail allowlist later)
        return raw

    return raw


def _prompt_target() -> str:
    """Ask the user for a URL / website / application name."""
    print("  Enter a target URL or website to scan.")
    print("  Examples:  example.com, https://app.example.com, My App (app.com)\n")
    sys.stdout.flush()
    while True:
        try:
            target_input = input("  > Target: ").strip()
            if not target_input:
                continue
            target = _normalize_target(target_input)
            if target and "." in target:
                return target
            print("  [!] Invalid input. Please enter a valid domain or URL.\n")
            sys.stdout.flush()
        except (EOFError, KeyboardInterrupt):
            sys.exit(0)
    return ""  # Satisfy linter


def _prompt_profile() -> str:
    """Let the user pick a scan intensity."""
    profiles = {
        "1": ("quick",    "Fast surface-level check (~2 min)"),
        "2": ("standard", "Balanced depth + speed  (~5 min)"),
        "3": ("deep",     "Full-depth comprehensive (~15 min)"),
        "4": ("ai_browser", "AI Browser Agent (Full interaction + Login support)"),
    }
    sys.stdout.write("\n  Select scan profile:\n\n")
    for key, (name, desc) in profiles.items():
        sys.stdout.write(f"    [{key}] {name:10s} — {desc}\n")
    sys.stdout.write("\n")
    sys.stdout.flush()
    while True:
        try:
            choice = input("  > Profile [1/2/3/4] (default 2): ").strip() or "2"
            if choice in profiles:
                p_item = profiles[choice]
                return p_item[0]
            print("  [!] Please enter 1, 2, 3, or 4.")
            sys.stdout.flush()
        except (EOFError, KeyboardInterrupt):
            sys.exit(0)
    return "standard"

def _prompt_auth() -> bool:
    """Ask whether to attempt authenticated scanning."""
    try:
        choice = input("  > Attempt authenticated scan? [y/N]: ").strip().lower()
        return choice in ("y", "yes")
    except (EOFError, KeyboardInterrupt):
        sys.exit(0)
    return False


def _auto_add_to_allowlist(config: dict, target: str) -> None:
    """Ensure the target domain is in the allowlist so safety checks pass."""
    enforcer = AllowlistEnforcer(config)
    if not enforcer.is_allowed(target):
        enforcer.add_domain(target)
        print(f"  \033[32m[+] Added '{target}' to allowlist.\033[0m")
    else:
        print(f"  \033[32m[OK] '{target}' is already in the allowlist.\033[0m")


# ── Post-scan Claude chat ─────────────────────────

_CHAT_BANNER = """
  \033[36m╔══════════════════════════════════════════════════╗
  ║        Claude AI — Post-Scan Assistant           ║
  ╚══════════════════════════════════════════════════╝\033[0m

  Ask questions about the scan findings, request remediation
  advice, explore attack chains, or get help prioritizing fixes.

  Commands:
    \033[33m/remediate <number>\033[0m  — Get detailed fix for finding #N
    \033[33m/summary\033[0m            — Regenerate the executive summary
    \033[33m/export\033[0m             — Save this chat to a file
    \033[33m/clear\033[0m              — Reset chat history
    \033[33m/quit\033[0m               — Exit the chat session
"""


async def _post_scan_chat(orchestrator: "Orchestrator", target: str) -> None:
    """Interactive post-scan Q&A session powered by Claude."""
    claude = orchestrator.claude

    # Ask if user wants to enter chat
    sys.stdout.write("  [*] Claude AI assistant is available for post-scan analysis.\n")
    sys.stdout.write("  > Open interactive chat? [Y/n]: ")
    sys.stdout.flush()
    choice = sys.stdin.readline().strip().lower()
    if choice in ("n", "no"):
        return

    print(_CHAT_BANNER)

    # Build scan context for Claude
    scan_context = {"results": orchestrator.results}

    # Initialize chat with scan data
    claude.reset_chat()

    # Track findings for /remediate command
    all_findings: list[dict] = []
    all_tech: list[dict] = []
    for res in orchestrator.results:
        all_findings.extend(res.get("vulnerabilities", []))
        all_tech.extend(res.get("tech_stack", []))

    # Show a quick findings summary to orient the user
    if all_findings:
        print(f"  \033[90mLoaded {len(all_findings)} findings from scan of {target}.\033[0m")
        sev_counts: dict[str, int] = {}
        for f in all_findings:
            s = f.get("severity", "info").lower()
            sev_counts[s] = sev_counts.get(s, 0) + 1
        sev_parts = [
            f"{count} {sev.upper()}" for sev, count in
            sorted(sev_counts.items(), key=lambda x: ["critical","high","medium","low","info"].index(x[0]) if x[0] in ["critical","high","medium","low","info"] else 99)
        ]
        print(f"  \033[90mBreakdown: {', '.join(sev_parts)}\033[0m\n")
    else:
        print(f"  \033[90mNo findings detected — you can still ask about hardening.\033[0m\n")

    chat_log: list[dict[str, str]] = []

    while True:
        try:
            sys.stdout.write("  you> ")
            sys.stdout.flush()
            user_input = sys.stdin.readline().strip()
        except Exception:
            sys.stdout.write("\n  Exiting chat.\n")
            sys.stdout.flush()
            break

        if not user_input:
            continue

        # ── Handle slash commands ──────────
        if user_input.lower() in ("/quit", "/exit", "/q"):
            print("  \033[90mExiting chat. Goodbye!\033[0m")
            break

        if user_input.lower() == "/clear":
            claude.reset_chat()
            chat_log.clear()
            print("  \033[90mChat history cleared.\033[0m\n")
            continue

        if user_input.lower() == "/export":
            await _export_chat_log(chat_log, target)
            continue

        if user_input.lower() == "/summary":
            user_input = (
                "Give me an executive summary of the scan results. "
                "Include the risk rating, top findings, and recommended priorities."
            )

        if user_input.lower().startswith("/remediate"):
            parts = user_input.split()
            if len(parts) >= 2 and parts[1].isdigit():
                idx = int(parts[1]) - 1
                if 0 <= idx < len(all_findings):
                    finding = all_findings[idx]
                    print(f"\n  \033[36m[*] Getting remediation for: {finding.get('id', '?')} "
                          f"({finding.get('severity', '?').upper()})\033[0m\n")
                    remediation = await claude.get_remediation(finding, all_tech)
                    if "[AUTH_REQUIRED]" in remediation:
                        print(f"  \033[33m{remediation}\033[0m")
                        input("  \033[33m>\033[0m Press Enter after authenticating... ")
                        remediation = await claude.get_remediation(finding, all_tech)
                    if remediation:
                        print(f"\n{_indent(remediation)}\n")
                        chat_log.append({"role": "user", "content": f"/remediate {parts[1]}"})
                        chat_log.append({"role": "assistant", "content": remediation})
                    else:
                        print("  \033[31m[!] Could not get remediation advice.\033[0m\n")
                else:
                    print(f"  \033[31m[!] Finding #{parts[1]} not found. Range: 1-{len(all_findings)}\033[0m\n")
            else:
                print("  \033[90mUsage: /remediate <number>  (e.g. /remediate 3)\033[0m\n")
            continue

        # ── Regular chat message ───────────
        print()  # breathing room

        reply = await claude.chat(user_input, scan_context=scan_context)

        # Handle Bonsai proxy auth redirect
        if "[AUTH_REQUIRED]" in reply:
            print(f"  \033[33m{reply}\033[0m")
            input("  \033[33m>\033[0m Press Enter after authenticating... ")
            # Retry the same message
            reply = await claude.chat(user_input, scan_context=scan_context)

        if reply:
            print(f"\n{_indent(reply)}\n")
            chat_log.append({"role": "user", "content": user_input})
            chat_log.append({"role": "assistant", "content": reply})
        else:
            print("  \033[31m[!] No response from Claude. Check connection.\033[0m\n")


def _indent(text: str, prefix: str = "  ") -> str:
    """Indent every line of text for clean terminal display."""
    return "\n".join(prefix + line for line in text.splitlines())


async def _export_chat_log(chat_log: list[dict[str, str]], target: str) -> None:
    """Save the chat session to a markdown file."""
    if not chat_log:
        print("  \033[90mNothing to export yet.\033[0m\n")
        return

    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    path = Path(f"./reports/chat_{target}_{timestamp}.md")
    path.parent.mkdir(parents=True, exist_ok=True)

    lines = [
        f"# Post-Scan Chat — {target}",
        f"**Date:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "",
        "---",
        "",
    ]
    for msg in chat_log:
        role = "You" if msg["role"] == "user" else "Claude"
        lines.append(f"### {role}")
        lines.append(msg["content"])
        lines.append("")

    path.write_text("\n".join(lines), encoding="utf-8")
    print(f"  \033[32m[+] Chat exported to {path}\033[0m\n")


# ── Async entry point ──────────────────────────────

async def interactive_main() -> None:
    """Interactive single-target scan flow."""
    _print_banner()

    # 1. Load config
    config_path = Path("config.yaml")
    if not config_path.exists():
        print("  \033[31m[!] config.yaml not found in current directory.\033[0m")
        sys.exit(1)
    config = load_config(config_path)

    # 2. Get target from user
    target = _prompt_target()
    print(f"\n  \033[36m[*] Target resolved: {target}\033[0m")

    # 3. Pick profile
    profile = _prompt_profile()
    print(f"  \033[36m[*] Profile: {profile}\033[0m")

    # 4. Auth toggle
    use_auth = _prompt_auth()
    config["auth"]["enabled"] = use_auth
    if use_auth:
        print("  \033[36m[*] Authenticated scan: ON\033[0m")
    else:
        print("  \033[36m[*] Authenticated scan: OFF\033[0m")

    # 5. Auto-authorize target in allowlist
    _auto_add_to_allowlist(config, target)

    # 6. Confirm and go
    sys.stdout.write(f"\n  {'=' * 50}\n")
    sys.stdout.write(f"    TARGET   :  {target}\n")
    sys.stdout.write(f"    PROFILE  :  {profile}\n")
    sys.stdout.write(f"    AUTH     :  {'yes' if use_auth else 'no'}\n")
    sys.stdout.write(f"  {'=' * 50}\n\n")
    sys.stdout.write("  > Start scan? [Y/n]: ")
    sys.stdout.flush()
    
    confirm_raw = sys.stdin.readline().strip().lower()
    if confirm_raw in ("n", "no"):
        sys.stdout.write("  Scan cancelled.\n")
        sys.stdout.flush()
        return

    print("\n  \033[32m[>] Launching scan...\033[0m\n")

    # 7. Build orchestrator and run
    orchestrator = Orchestrator(
        config=config,
        targets=[target],
        profile=profile,
    )

    # Graceful shutdown
    import platform
    if platform.system() != "Windows":
        loop = asyncio.get_running_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            loop.add_signal_handler(
                sig,
                lambda s=sig: asyncio.create_task(orchestrator.shutdown(s)),
            )
    else:
        signal.signal(
            signal.SIGINT,
            lambda s, f: asyncio.create_task(orchestrator.shutdown(signal.SIGINT)),
        )

    await orchestrator.run()

    print("\n  \033[32m[OK] Scan complete! Reports saved to ./reports/\033[0m\n")

    # ── Post-scan interactive Claude chat ──────────
    if orchestrator.claude.available:
        await _post_scan_chat(orchestrator, target)


async def batch_main(args: argparse.Namespace) -> None:
    """Legacy batch mode — reads targets from a file."""
    config = load_config(args.config)
    if args.output:
        config["reporting"]["output_dir"] = str(args.output)

    if not args.targets.exists():
        log.error("targets_file_missing", path=str(args.targets))
        sys.exit(1)

    targets = [
        line.strip()
        for line in args.targets.read_text().splitlines()
        if line.strip() and not line.strip().startswith("#")
    ]

    if not targets:
        log.error("no_targets_provided")
        sys.exit(1)

    log.info(
        "orchestrator_starting",
        target_count=len(targets),
        profile=args.profile,
    )

    orchestrator = Orchestrator(
        config=config,
        targets=targets,
        profile=args.profile,
    )

    import platform
    if platform.system() != "Windows":
        loop = asyncio.get_running_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            loop.add_signal_handler(
                sig,
                lambda s=sig: asyncio.create_task(orchestrator.shutdown(s)),
            )
    else:
        signal.signal(
            signal.SIGINT,
            lambda s, f: asyncio.create_task(orchestrator.shutdown(signal.SIGINT)),
        )

    await orchestrator.run()


# ── CLI argument parsing ───────────────────────────

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Automated Vulnerability Scanning Orchestrator"
    )
    p.add_argument(
        "--batch", action="store_true",
        help="Run in batch mode (read targets from file instead of interactive prompt)"
    )
    p.add_argument(
        "--config", type=Path, default=Path("config.yaml"),
        help="Path to config.yaml"
    )
    p.add_argument(
        "--targets", type=Path, default=Path("targets.txt"),
        help="File with target domains (batch mode only)"
    )
    p.add_argument(
        "--profile", choices=["quick", "standard", "deep"],
        default="standard", help="Scan intensity profile (batch mode only)"
    )
    p.add_argument(
        "--output", type=Path, default=None,
        help="Override report output directory"
    )
    return p.parse_args()


def main() -> None:
    structlog.configure(
        processors=[
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.add_log_level,
            structlog.dev.ConsoleRenderer(),
        ],
        wrapper_class=structlog.make_filtering_bound_logger(
            __import__("logging").INFO
        ),
    )

    args = parse_args()
    
    # Always run bonsai login check if config says it's enabled
    config = load_config(args.config)
    if config.get("claude_assistant", {}).get("enabled", False) or config.get("ai_analysis", {}).get("enabled", False):
        run_bonsai_login()

    if args.batch:
        asyncio.run(batch_main(args))
    else:
        asyncio.run(interactive_main())


if __name__ == "__main__":
    main()