main.py

#!/usr/bin/env python3
"""
main.py — user-friendly entry point for the LLM-Powered Pentest Agent.

Just run:
    python3 main.py

The agent will walk you through:
  1. Authorization check
  2. Target URL
  3. Test mode (black box / gray box)
  4. Credentials (if gray box)
  5. Scan guidance / suggestions (optional, multi-line)
  6. Scan start
  7. Single .xlsx report + summary sheet saved to ./reports/

CLI flags are still available for power users / scripting, but everything
has sensible defaults so the basic case needs zero flags.

A `--quick` flag accepts all defaults and only asks for the target URL,
suitable for one-liner scripted use.
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from datetime import datetime
from typing import Any, Dict, List, Optional

from colorama import Fore, Style, init as _colorama_init

from config import load_config
from core.scan_suggestions import (
    ScanSuggestions,
    parse_suggestions_from_string,
)
from pentest_agent import PentestAgent

try:
    import openpyxl
    HAVE_OPENPYXL = True
except Exception:
    HAVE_OPENPYXL = False

# Initialize colorama (auto-disables on non-tty)
_colorama_init(autoreset=True)


# ============================================================== banner


BANNER = r"""
==============================================================
  LLM-Powered Deep Penetration Testing Agent  v1.0
  Plan-aligned orchestrator + 46 scanners / 11 categories
  https://rashedtech.com/
==============================================================
"""


# ============================================================== helpers


def _say(msg: str, color: str = "") -> None:
    if color and sys.stdout.isatty():
        print(color + msg + Style.RESET_ALL)
    else:
        print(msg)


def _prompt(label: str, *, default: str = "", secret: bool = False) -> str:
    suffix = f" [{default}]" if default else ""
    while True:
        try:
            if secret:
                import getpass
                val = getpass.getpass(f"  {label}{suffix}: ").strip()
            else:
                val = input(f"  {label}{suffix}: ").strip()
            if not val and default:
                return default
            return val
        except (KeyboardInterrupt, EOFError):
            print()
            _say("\nAborted.", Fore.YELLOW)
            sys.exit(130)


def _confirm_authorization() -> None:
    print(BANNER)
    _say(Fore.RED + "AUTHORIZATION CHECK" + Style.RESET_ALL)
    print()
    print("  This tool actively probes (and may trigger) real vulnerabilities.")
    print("  Only use it against a target you OWN or are EXPLICITLY AUTHORIZED")
    print("  to test (your own app, a signed engagement, a bug bounty in-scope")
    print("  asset, or a legal practice target like OWASP Juice Shop / DVWA).")
    print()
    print('  Type exactly:  I CONFIRM I AM AUTHORIZED')
    print()
    while True:
        ans = input("  > ").strip()
        if ans == "I CONFIRM I AM AUTHORIZED":
            _say("  ✓ Authorization confirmed.", Fore.GREEN)
            print()
            return
        _say("  ✗ Type the phrase exactly as shown.", Fore.RED)
        if ans.lower() in ("exit", "quit", "q"):
            sys.exit(0)


def _ask_target() -> str:
    print("STEP 1 of 5 — TARGET")
    print("  Enter the URL to test. Include the scheme (https://…).")
    while True:
        url = _prompt("Target URL").strip()
        if not url:
            continue
        if not url.startswith("http://") and not url.startswith("https://"):
            url = "https://" + url
        url = url.rstrip("/")
        # Friendly confirmation
        print()
        _say(f"  → Target set to: {Fore.CYAN}{url}{Style.RESET_ALL}")
        ok = _prompt("Confirm? [Y/n]", default="Y").lower()
        if ok in ("", "y", "yes"):
            print()
            return url
        # else: loop and re-ask


def _ask_mode() -> str:
    print("STEP 2 of 5 — TEST MODE")
    print()
    print("  [1] Black Box    — no credentials, anonymous testing")
    print("  [2] Gray Box     — you have credentials, authenticated testing")
    print()
    while True:
        choice = _prompt("Choose [1/2]", default="1").strip()
        if choice in ("1", ""):
            print()
            _say(f"  → Mode: {Fore.CYAN}black box{Style.RESET_ALL}")
            print()
            return "black_box"
        if choice == "2":
            print()
            _say(f"  → Mode: {Fore.CYAN}gray box{Style.RESET_ALL}")
            print()
            return "gray_box"
        _say("  Please enter 1 or 2.", Fore.YELLOW)


def _ask_credentials() -> Optional[Dict[str, Any]]:
    print("STEP 3 of 5 — CREDENTIALS (gray box)")
    print()
    print("  Provide ONE of:")
    print("    • username + password        (we'll try common login shapes)")
    print("    • bearer / API key          (Authorization: Bearer …)")
    print("    • leave blank                (we'll fall back to black box)")
    print()
    username = _prompt("Username (blank to skip)").strip()
    if not username:
        print()
        _say("  → No credentials — falling back to black box.", Fore.YELLOW)
        print()
        return None
    password = _prompt("Password", secret=True)
    bearer = _prompt("Bearer / API key (optional)").strip()
    extras: Dict[str, str] = {}
    if bearer:
        extras["bearer"] = bearer
    print()
    _say(f"  → Credentials captured for user: {Fore.CYAN}{username}{Style.RESET_ALL}", Fore.GREEN)
    print()
    return {"username": username, "password": password, "extra": extras}


def _ask_suggestions() -> ScanSuggestions:
    print("STEP 4 of 5 — SCAN GUIDANCE (optional)")
    print()
    print("  Tell the agent how to scan. One directive per line.")
    print("  Examples:")
    _say(f"    {Fore.CYAN}ignore path: /logout, /api/health{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}ignore subdomain: staging., dev.{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}focus path: /api, /admin{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}user agent: Mozilla/5.0 ...{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}header X-Engagement: pentest-2026-q1{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}delay: 0.5{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}max requests: 2000{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}include only: sql_injection, xss_testing, hardcoded_secrets{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}exclude: prompt_injection, model_extraction{Style.RESET_ALL}")
    _say(f"    {Fore.CYAN}notes: This is a Q4 re-test. Focus on /api endpoints.{Style.RESET_ALL}")
    print()
    print("  Type your directives, then press ENTER on an empty line to continue.")
    print("  (or just press ENTER to skip and use defaults)")
    print()
    lines: List[str] = []
    while True:
        try:
            line = input("  > ")
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not line.strip():
            break
        lines.append(line)
    suggestions = parse_suggestions_from_string("\n".join(lines))
    if not suggestions.is_empty():
        print()
        _say("  ✓ Scan guidance captured:", Fore.GREEN)
        for k, v in suggestions.to_dict().items():
            if v and v != [] and v != {} and v != "" and k != "rotate_user_agent":
                _say(f"      • {k} = {v}", Fore.CYAN)
    else:
        print()
        _say("  → No guidance — using defaults.", Fore.YELLOW)
    print()
    return suggestions


def _confirm_and_run(target: str, mode: str, creds: Optional[Dict[str, Any]],
                     suggestions: ScanSuggestions, cfg: Dict[str, Any],
                     auto_start: bool = False) -> Dict[str, Any]:
    print("STEP 5 of 5 — READY TO SCAN")
    print()
    _say(f"  Target      : {Fore.CYAN}{target}{Style.RESET_ALL}")
    print(f"  Mode        : {mode}")
    if creds:
        print(f"  User        : {creds.get('username')}")
    if not suggestions.is_empty():
        n = sum(1 for v in suggestions.to_dict().values() if v)
        print(f"  Guidance    : {n} directive(s)")
    print()
    if auto_start:
        _say("  ✓ Auto-start (--yes)", Fore.GREEN)
    else:
        ans = _prompt("Start scan? [Y/n]", default="Y").lower()
        if ans not in ("", "y", "yes"):
            _say("Aborted.", Fore.YELLOW)
            sys.exit(0)
    print()
    print("=" * 60)
    print("SCAN STARTING")
    print("=" * 60)
    print()

    agent = PentestAgent(cfg)
    started = time.time()
    result = agent.run_pentest(
        target_url=target,
        test_mode=mode,
        credentials=creds,
        output_path="reports",
        output_format=["excel", "summary"],
        log=_progress_log(),
        suggestions=suggestions,
    )
    elapsed = time.time() - started
    return {**result, "elapsed_seconds": elapsed}


def _progress_log():
    """Build a log callable that prints short, color-coded progress lines."""
    started = time.time()

    def log(msg: str) -> None:
        # Strip leading newline so we can prepend our own timestamp
        msg = msg.strip()
        if not msg:
            return
        elapsed = int(time.time() - started)
        # Color phase banners
        if msg.startswith("PHASE"):
            _say(f"[{elapsed:>4}s] {Fore.MAGENTA}{msg}{Style.RESET_ALL}")
        elif msg.startswith("==="):
            _say(f"[{elapsed:>4}s] {Fore.MAGENTA}{msg}{Style.RESET_ALL}")
        elif "FINDING" in msg.upper() or "🔴" in msg:
            _say(f"[{elapsed:>4}s] {Fore.RED}{msg}{Style.RESET_ALL}")
        elif "🟢" in msg or "tested, nothing found" in msg:
            _say(f"[{elapsed:>4}s] {Fore.GREEN}{msg}{Style.RESET_ALL}")
        elif "FAIL" in msg.upper() or "⚠️" in msg or "crashed" in msg.lower():
            _say(f"[{elapsed:>4}s] {Fore.YELLOW}{msg}{Style.RESET_ALL}")
        else:
            print(f"[{elapsed:>4}s] {msg}", flush=True)

    return log


def _show_result(result: Dict[str, Any]) -> None:
    print()
    print("=" * 60)
    _say(Fore.GREEN + "SCAN COMPLETE" + Style.RESET_ALL)
    print("=" * 60)
    print()
    findings = result.get("findings", [])
    # Severity breakdown
    from collections import Counter
    sev = Counter(f.get("severity", "info").lower() for f in findings)
    print(f"  Time elapsed     : {result.get('elapsed_seconds', 0):.1f}s")
    print(f"  Total findings   : {len(findings)}")
    for s in ("critical", "high", "medium", "low", "info"):
        if sev.get(s):
            color = {"critical": Fore.RED, "high": Fore.RED,
                     "medium": Fore.YELLOW, "low": Fore.CYAN, "info": ""}[s]
            _say(f"    {color}{s.capitalize():<10}: {sev[s]}{Style.RESET_ALL}")
    print()
    print("  REPORTS")
    for k, v in result.get("report_paths", {}).items():
        if k.endswith("_size_kb"):
            continue
        marker = "📊" if k == "excel" else ("📄" if k == "summary" else "📁")
        size = result["report_paths"].get(f"{k}_size_kb", "")
        size_s = f"  ({size})" if size else ""
        _say(f"    {marker} {Fore.CYAN}{v}{Style.RESET_ALL}{size_s}")
    print()
    print("  WHAT TO DO NEXT")
    print("    1. Open the .xlsx — every finding has remediation steps.")
    print("    2. Pivot on the 'Scanner Coverage' sheet — nothing was skipped.")
    print("    3. Re-test after fixes with the same suggestions file.")
    print()


def _open_report(path: str) -> None:
    """Best-effort open the xlsx in the user's default app."""
    if not HAVE_OPENPYXL:
        return
    try:
        if sys.platform.startswith("darwin"):
            os.system(f'open "{path}"')
        elif os.name == "nt":
            os.startfile(path)  # type: ignore
        elif os.name == "posix":
            # xdg-open is on virtually every Linux desktop
            os.system(f'xdg-open "{path}" >/dev/null 2>&1 &')
    except Exception:
        pass


# ============================================================== entry


def main() -> int:
    parser = argparse.ArgumentParser(
        description="LLM-Powered Pentest Agent — user-friendly mode.",
        add_help=True,
    )
    # Power-user overrides. Everything is optional; default is the friendly flow.
    parser.add_argument("--target", help="Skip prompt, use this target URL")
    parser.add_argument("--mode", choices=["black_box", "gray_box"], help="Skip prompt")
    parser.add_argument("--username", help="Skip prompt (gray_box)")
    parser.add_argument("--password", help="Skip prompt (gray_box)")
    parser.add_argument("--bearer", help="Skip prompt (gray_box)")
    parser.add_argument("--suggestions-file", help="Read guidance from a text file")
    parser.add_argument("--output", default="reports", help="Output directory")
    parser.add_argument("--yes", "-y", action="store_true",
                        help="Skip authorization confirmation (for scripted runs)")
    parser.add_argument("--quick", action="store_true",
                        help="Accept all defaults; only ask for the target URL")
    parser.add_argument("--no-open", action="store_true",
                        help="Don't auto-open the xlsx when the scan finishes")
    parser.add_argument("--email", action="store_true",
                        help="Email the xlsx (requires config.email.*)")
    args = parser.parse_args()

    # ---- 1. Authorization ----
    if not args.yes and not args.target:
        _confirm_authorization()

    # ---- 2. Target ----
    if args.target:
        target = args.target.strip()
        if not target.startswith("http://") and not target.startswith("https://"):
            target = "https://" + target
        target = target.rstrip("/")
    else:
        target = _ask_target()

    # ---- 3. Mode ----
    if args.mode:
        mode = args.mode
    elif args.username or args.password or args.bearer:
        mode = "gray_box"
    elif args.quick:
        mode = "black_box"
    else:
        mode = _ask_mode()

    # ---- 4. Credentials (gray box) ----
    creds: Optional[Dict[str, Any]] = None
    if mode == "gray_box":
        if args.username or args.password or args.bearer:
            creds = {
                "username": args.username or "",
                "password": args.password or "",
                "extra": {"bearer": args.bearer} if args.bearer else {},
            }
        elif args.target and args.yes:
            # Non-interactive gray_box without credentials: skip silently
            mode = "black_box"
        else:
            creds = _ask_credentials()
            if creds is None:
                mode = "black_box"

    # ---- 5. Suggestions ----
    if args.suggestions_file:
        try:
            with open(args.suggestions_file, "r", encoding="utf-8") as fh:
                suggestions = parse_suggestions_from_string(fh.read())
        except FileNotFoundError:
            _say(f"  Suggestions file not found: {args.suggestions_file}", Fore.YELLOW)
            suggestions = ScanSuggestions()
    elif args.quick or (args.target and args.yes):
        # Non-interactive path: no suggestions
        suggestions = ScanSuggestions()
    else:
        suggestions = _ask_suggestions()

    # ---- Run ----
    cfg = load_config()
    auto_start = bool(args.yes)
    result = _confirm_and_run(target, mode, creds, suggestions, cfg, auto_start=auto_start)

    # ---- Optional email delivery ----
    if args.email:
        excel_path = result.get("report_paths", {}).get("excel")
        if excel_path:
            try:
                from reporting.delivery import email_xlsx
                ok, info = email_xlsx(excel_path, cfg.get("email", {}),
                                      subject=f"Pentest report — {target}")
                if ok:
                    _say(f"  ✓ Emailed: {info}", Fore.GREEN)
                else:
                    _say(f"  ✗ Email failed: {info}", Fore.YELLOW)
            except Exception as e:
                _say(f"  ✗ Email failed: {e}", Fore.YELLOW)

    # ---- Show results + open ----
    _show_result(result)
    excel_path = result.get("report_paths", {}).get("excel")
    if excel_path and not args.no_open and sys.stdout.isatty():
        _say("Opening the report in your default app…", Fore.CYAN)
        _open_report(excel_path)
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        _say("\nAborted.", Fore.YELLOW)
        sys.exit(130)