cli.py

#!/usr/bin/env python3
"""
DPI-Bench CLI — command-line interface for DPI evasion testing.

Usage:
    dpi-bench scan <target> [--config <path>] [--timeout <sec>]
    dpi-bench test <target> --technique <name> [--config <path>]
    dpi-bench bench <target> [--config <path>] [--output <dir>]
    dpi-bench report <results.json> [--format <fmt>]
    dpi-bench detect <target> [--timeout <sec>]
    dpi-bench list-techniques
"""

import json
import logging
import sys
from pathlib import Path
from typing import Optional

import click

from dpibench import __version__
from dpibench.core.config import Config, ConfigError
from dpibench.core.engine import TestEngine
from dpibench.reporting.console import ConsoleReporter
from dpibench.reporting.html_report import HTMLReporter
from dpibench.reporting.json_export import JSONExporter, CSVExporter
from dpibench.targets.probes import DPIProber

logger = logging.getLogger("dpibench")


def setup_logging(verbose: bool = False, quiet: bool = False):
    """Configure logging based on verbosity flags."""
    if quiet:
        level = logging.WARNING
    elif verbose:
        level = logging.DEBUG
    else:
        level = logging.INFO

    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%H:%M:%S",
    )


def load_config(config_path: Optional[str], targets_path: Optional[str] = None) -> Config:
    """Load configuration from file or defaults."""
    try:
        if config_path:
            cfg = Config(config_path)
        else:
            default_path = Path("config/default.yaml")
            if default_path.exists():
                cfg = Config(str(default_path))
            else:
                cfg = Config()

        if targets_path:
            cfg.load_targets_file(targets_path)

        return cfg
    except ConfigError as e:
        click.echo(f"Configuration error: {e}", err=True)
        sys.exit(1)


@click.group()
@click.version_option(version=__version__, prog_name="dpi-bench")
@click.option("-v", "--verbose", is_flag=True, help="Enable verbose output")
@click.option("-q", "--quiet", is_flag=True, help="Suppress non-error output")
@click.pass_context
def main(ctx, verbose, quiet):
    """DPI-Bench: DPI evasion testing and benchmarking toolkit."""
    setup_logging(verbose, quiet)
    ctx.ensure_object(dict)
    ctx.obj["verbose"] = verbose
    ctx.obj["quiet"] = quiet


@main.command()
@click.argument("target")
@click.option("-c", "--config", "config_path", default=None, help="Config file path")
@click.option("-t", "--targets", "targets_path", default=None, help="Targets file path")
@click.option("--timeout", default=10.0, help="Connection timeout in seconds")
@click.option("--no-validate", is_flag=True, help="Skip target validation")
@click.option("--no-detect", is_flag=True, help="Skip DPI detection")
@click.pass_context
def scan(ctx, target, config_path, targets_path, timeout, no_validate, no_detect):
    """Run a full scan against a target domain."""
    cfg = load_config(config_path, targets_path)

    # Add CLI target if not in config
    if not cfg.targets:
        from dpibench.core.config import TargetProfile
        tp = TargetProfile("cli", [{"domain": target, "port": 443, "protocol": "https", "sni": target}])
        cfg._targets.append(tp)

    cfg._connection["timeout"] = timeout

    engine = TestEngine(config=cfg)
    suite_result = engine.run(skip_validation=no_validate, skip_dpi_detect=no_detect)

    # Console report
    reporter = ConsoleReporter(verbose=ctx.obj["verbose"])
    reporter.report(suite_result.to_dict())


@main.command()
@click.argument("target")
@click.option("--technique", "-T", required=True, help="Technique name to test")
@click.option("-c", "--config", "config_path", default=None, help="Config file path")
@click.option("--timeout", default=10.0, help="Connection timeout in seconds")
@click.option("--port", default=443, help="Target port")
@click.option("--sni", default=None, help="Override SNI value")
@click.pass_context
def test(ctx, target, technique, config_path, timeout, port, sni):
    """Test a specific technique against a target."""
    cfg = load_config(config_path)
    cfg._connection["timeout"] = timeout

    engine = TestEngine(config=cfg)

    available = engine.list_techniques()
    if technique not in available:
        click.echo(f"Unknown technique: {technique}")
        click.echo(f"Available: {', '.join(available)}")
        sys.exit(1)

    click.echo(f"Testing {technique} against {target}:{port}...")
    results = engine.run_single_technique(technique, target, port, sni)

    reporter = ConsoleReporter(verbose=ctx.obj["verbose"])
    reporter.print_header(f"Test: {technique}")
    reporter.print_results_table([r.to_dict() for r in results])

    passed = sum(1 for r in results if r.success)
    click.echo(f"\nResult: {passed}/{len(results)} passed")


@main.command()
@click.argument("target")
@click.option("-c", "--config", "config_path", default=None, help="Config file path")
@click.option("-t", "--targets", "targets_path", default=None, help="Targets file path")
@click.option("-o", "--output", "output_dir", default="reports/", help="Output directory")
@click.option("--timeout", default=10.0, help="Connection timeout in seconds")
@click.option("--format", "formats", multiple=True, default=["console", "json"],
              help="Output formats (console, json, csv, html)")
@click.pass_context
def bench(ctx, target, config_path, targets_path, output_dir, timeout, formats):
    """Run a full benchmark suite with reports."""
    cfg = load_config(config_path, targets_path)

    if not cfg.targets:
        from dpibench.core.config import TargetProfile
        tp = TargetProfile("cli", [{"domain": target, "port": 443, "protocol": "https", "sni": target}])
        cfg._targets.append(tp)

    cfg._connection["timeout"] = timeout
    engine = TestEngine(config=cfg)

    click.echo(f"Running benchmark against {target}...")
    suite_result = engine.run()
    result_dict = suite_result.to_dict()

    for fmt in formats:
        if fmt == "console":
            reporter = ConsoleReporter(verbose=ctx.obj["verbose"])
            reporter.report(result_dict)
        elif fmt == "json":
            exporter = JSONExporter(output_dir=output_dir)
            path = exporter.export(result_dict)
            click.echo(f"JSON report: {path}")
        elif fmt == "csv":
            exporter = CSVExporter(output_dir=output_dir)
            path = exporter.export(result_dict)
            click.echo(f"CSV report: {path}")
            path = exporter.export_matrix(result_dict)
            click.echo(f"CSV matrix: {path}")
        elif fmt == "html":
            exporter = HTMLReporter(output_dir=output_dir)
            path = exporter.generate(result_dict)
            click.echo(f"HTML report: {path}")


@main.command()
@click.argument("results_file")
@click.option("--format", "fmt", default="console", help="Output format (console, html)")
@click.option("-o", "--output", "output_dir", default="reports/", help="Output directory")
@click.pass_context
def report(ctx, results_file, fmt, output_dir):
    """Generate a report from saved JSON results."""
    results_path = Path(results_file)
    if not results_path.exists():
        click.echo(f"Results file not found: {results_file}", err=True)
        sys.exit(1)

    try:
        with open(results_path, "r", encoding="utf-8") as f:
            data = json.load(f)
    except (json.JSONDecodeError, OSError) as e:
        click.echo(f"Failed to load results: {e}", err=True)
        sys.exit(1)

    if fmt == "console":
        reporter = ConsoleReporter(verbose=ctx.obj["verbose"])
        reporter.report(data)
    elif fmt == "html":
        exporter = HTMLReporter(output_dir=output_dir)
        path = exporter.generate(data)
        click.echo(f"HTML report: {path}")
    else:
        click.echo(f"Unknown format: {fmt}", err=True)
        sys.exit(1)


@main.command()
@click.argument("target")
@click.option("--timeout", default=10.0, help="Connection timeout in seconds")
@click.pass_context
def detect(ctx, target, timeout):
    """Detect the DPI system on the path to a target."""
    click.echo(f"Probing DPI on path to {target}...")
    prober = DPIProber(timeout=timeout)
    result = prober.identify_dpi(target)

    reporter = ConsoleReporter(verbose=ctx.obj["verbose"])
    reporter.print_header("DPI Detection")
    reporter.print_dpi_info(result.to_dict())


@main.command("list-techniques")
def list_techniques():
    """List all available evasion techniques."""
    cfg = Config()
    engine = TestEngine(config=cfg)
    techniques = engine.list_techniques()

    click.echo("Available techniques:")
    for name in sorted(techniques):
        info = engine.get_technique_info(name)
        doc = (info.get("doc") or "").strip().split("\n")[0] if info else ""
        click.echo(f"  {name:<20} {doc}")


if __name__ == "__main__":
    main()