pentest.py

#!/usr/bin/env python3
"""
inmydata PenTest - Automated penetration testing for agentic-built web apps.

Orchestrates OWASP ZAP and Nuclei scanning with AWS Cognito authentication
and unified HTML reporting.
"""

import sys
import click
import yaml
from pathlib import Path
from datetime import datetime
from rich.console import Console

from src.orchestrator import PenTestOrchestrator

console = Console()


def load_config(config_path: str) -> dict:
    """Load and validate configuration from YAML file."""
    path = Path(config_path)
    if not path.exists():
        console.print(f"[red]Config file not found: {config_path}[/red]")
        console.print("Run: cp config.example.yaml config.yaml")
        sys.exit(1)

    with open(path) as f:
        config = yaml.safe_load(f)

    return config


@click.command()
@click.option(
    "--config", "-c",
    default="config.yaml",
    help="Path to configuration file"
)
@click.option(
    "--target", "-t",
    default=None,
    help="Target URL (overrides config file)"
)
@click.option(
    "--mode", "-m",
    type=click.Choice(["full", "quick", "nuclei-only"]),
    default=None,
    help="Scan mode (overrides config file)"
)
@click.option(
    "--output", "-o",
    default=None,
    help="Output directory for reports"
)
@click.option(
    "--no-auth",
    is_flag=True,
    default=False,
    help="Skip authentication (scan public endpoints only)"
)
@click.option(
    "--dry-run",
    is_flag=True,
    default=False,
    help="Validate config and show what would run without scanning"
)
def main(config, target, mode, output, no_auth, dry_run):
    """inmydata PenTest - Automated security scanning for web applications."""

    console.print()
    console.print("[bold blue]═══════════════════════════════════════════[/bold blue]")
    console.print("[bold blue]  inmydata PenTest[/bold blue]")
    console.print("[bold blue]  Automated Penetration Testing Pipeline[/bold blue]")
    console.print("[bold blue]═══════════════════════════════════════════[/bold blue]")
    console.print()

    # Load configuration
    cfg = load_config(config)

    # Apply CLI overrides
    if target:
        cfg["target"]["url"] = target
    if mode:
        cfg["scan"]["mode"] = mode
    if output:
        cfg["report"]["output_dir"] = output
    if no_auth:
        cfg["auth"]["provider"] = "none"

    target_url = cfg["target"]["url"]
    scan_mode = cfg["scan"]["mode"]

    console.print(f"[cyan]Target:[/cyan]     {target_url}")
    console.print(f"[cyan]Mode:[/cyan]       {scan_mode}")
    console.print(f"[cyan]Auth:[/cyan]       {cfg['auth']['provider']}")
    console.print(f"[cyan]Started:[/cyan]    {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    console.print()

    if dry_run:
        console.print("[yellow]DRY RUN - showing scan plan without executing[/yellow]")
        console.print()
        _show_scan_plan(cfg)
        return

    # Run the orchestrator
    orchestrator = PenTestOrchestrator(cfg)

    try:
        reports = orchestrator.run()

        console.print()
        console.print("[bold green]═══════════════════════════════════════════[/bold green]")
        console.print("[bold green]  Scan Complete[/bold green]")
        console.print("[bold green]═══════════════════════════════════════════[/bold green]")
        console.print()
        console.print(f"[green]HTML report:     {reports['html']}[/green]")
        console.print(f"[green]Markdown report: {reports['markdown']}[/green]")

        # CI/CD exit code logic
        if orchestrator.should_fail_pipeline():
            console.print(
                f"[red]PIPELINE FAIL: Findings at or above "
                f"'{cfg['ci']['fail_on_severity']}' severity detected[/red]"
            )
            sys.exit(1)
        else:
            console.print("[green]PIPELINE PASS: No blocking findings[/green]")
            sys.exit(0)

    except KeyboardInterrupt:
        console.print("\n[yellow]Scan interrupted by user[/yellow]")
        orchestrator.cleanup()
        sys.exit(130)
    except Exception as e:
        console.print(f"\n[red]Scan failed: {e}[/red]")
        orchestrator.cleanup()
        sys.exit(2)


def _show_scan_plan(cfg: dict):
    """Display what the scan would do without executing."""
    mode = cfg["scan"]["mode"]

    console.print("[bold]Scan Plan:[/bold]")
    console.print()

    if mode in ("full", "quick"):
        console.print("  [cyan]1. Authentication[/cyan]")
        if cfg["auth"]["provider"] == "cognito":
            console.print(f"     Authenticate via AWS Cognito ({cfg['auth']['cognito']['region']})")
        else:
            console.print("     Skipped (no-auth mode)")

        console.print("  [cyan]2. OWASP ZAP Spider[/cyan]")
        console.print(f"     Max depth: {cfg['zap']['spider']['max_depth']}")
        console.print(f"     Ajax spider: {cfg['zap']['spider']['ajax_spider']}")

        console.print("  [cyan]3. OWASP ZAP Passive Scan[/cyan]")
        console.print("     Runs automatically during spidering")

        if mode == "full":
            console.print("  [cyan]4. OWASP ZAP Active Scan[/cyan]")
            console.print(f"     Policy: {cfg['zap']['active_scan']['policy']}")
            console.print(f"     Threads: {cfg['zap']['active_scan']['threads']}")

    if cfg["nuclei"]["enabled"]:
        step = "5" if mode == "full" else ("4" if mode == "quick" else "1")
        console.print(f"  [cyan]{step}. Nuclei Template Scan[/cyan]")
        console.print(f"     Community templates: {cfg['nuclei']['templates']['community']}")
        console.print(f"     Custom templates: {cfg['nuclei']['templates']['custom_path']}")
        console.print(f"     Severity filter: {cfg['nuclei']['templates']['community_severity']}")

    console.print()
    console.print(f"  [cyan]Report:[/cyan] HTML + Markdown to {cfg['report']['output_dir']}/")
    console.print(f"  [cyan]CI gate:[/cyan] Fail on {cfg['ci']['fail_on_severity']}+ severity")


if __name__ == "__main__":
    main()