cli.py

#!/usr/bin/env python3
"""
LLMRed CLI

Usage:
    python cli.py --config config.yaml
    python cli.py --config config.yaml --nist-answers answers.yaml
    python cli.py --config config.yaml --dry-run
    python cli.py --config config_a.yaml --compare-config config_b.yaml
"""
import argparse
import json
import logging
import sys
from pathlib import Path

import yaml

from pentest.comparison import compare
from pentest.config import ConfigError, load_config
from pentest.core import Orchestrator
from pentest.reporting.comparison_generator import ComparisonReportGenerator
from pentest.policy import PolicyViolation
from pentest.validation import ReviewFileError, apply_human_reviews, load_review_file

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("cli")


def _load_questionnaire(path):
    if path is None:
        return None
    with open(path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    if not isinstance(data, dict):
        raise ValueError("--nist-answers must contain a YAML mapping.")
    return data


def _resolve_nist_answers(config, config_path, explicit_path=None):
    """Resolve NIST evidence without changing the main assessment command.

    Precedence: explicit CLI override, configured audit.nist_answers_file,
    then conventional nist-answers.yaml beside the main config. If NIST
    assessment is enabled but no evidence file exists, the report still maps
    technical findings and marks organizational outcomes Not assessed.
    """
    if explicit_path:
        return Path(explicit_path)
    if not config.audit.nist_questionnaire:
        return None

    config_dir = Path(config_path).expanduser().resolve().parent
    configured = config.audit.nist_answers_file
    candidate = Path(configured).expanduser() if configured else config_dir / "nist-answers.yaml"
    if configured and not candidate.is_absolute():
        candidate = config_dir / candidate
    if candidate.exists():
        logger.info("Loading integrated NIST AI RMF evidence: %s", candidate)
        return candidate
    if configured:
        raise ValueError(
            f"audit.nist_answers_file points to missing file: {candidate}"
        )
    logger.info(
        "No nist-answers.yaml found beside %s; technical NIST mappings will be "
        "reported and organizational outcomes will remain Not assessed.",
        config_path,
    )
    return None


def _dry_run(orchestrator, label: str) -> bool:
    resp = orchestrator.client.chat([{"role": "user", "content": "Reply with OK."}])
    if resp and resp.text:
        logger.info("[%s] Connectivity check succeeded. Response: %s", label, resp.text[:100])
        return True
    logger.error("[%s] Connectivity check failed — no usable reply text extracted.", label)
    if resp is not None and resp.raw:
        # Text extraction failed but a response DID come back - this is the
        # single most useful line for debugging a new/custom client target,
        # since it shows exactly what the server actually returned.
        logger.error("[%s] Raw response for debugging: %s", label, json.dumps(resp.raw)[:800])
    return False


def run_single(config, args):
    orchestrator = Orchestrator(config)

    if args.dry_run:
        ok = _dry_run(orchestrator, config.engagement.client_name)
        sys.exit(0 if ok else 2)

    findings = orchestrator.run()
    logger.info("Assessment complete: %d findings.", len(findings))

    if args.review_findings:
        review_data = load_review_file(Path(args.review_findings))
        orchestrator.findings = apply_human_reviews(orchestrator.findings, review_data)
        findings = orchestrator.findings
        logger.info("Applied %d human review decision(s).", len(review_data["reviews"]))

    questionnaire_answers = _load_questionnaire(
        _resolve_nist_answers(config, args.config, args.nist_answers)
    )
    nist = orchestrator.score_nist(questionnaire_answers)
    html_path, pdf_path, json_path = orchestrator.generate_report(nist)

    print(f"\nReport: {html_path}")
    if pdf_path:
        print(f"PDF:    {pdf_path}")
    if json_path:
        print(f"JSON:   {json_path}")
    if orchestrator.forensic_manifest_path:
        print(f"Evidence manifest: {orchestrator.forensic_manifest_path}")


def run_comparison(config_a, config_b, args):
    orchestrator_a = Orchestrator(config_a)
    orchestrator_b = Orchestrator(config_b)

    if args.dry_run:
        ok_a = _dry_run(orchestrator_a, config_a.engagement.client_name)
        ok_b = _dry_run(orchestrator_b, config_b.engagement.client_name)
        sys.exit(0 if (ok_a and ok_b) else 2)

    logger.info("Running full suite against target A: %s", config_a.engagement.client_name)
    orchestrator_a.run()
    logger.info("Target A complete: %d findings.", len(orchestrator_a.findings))

    logger.info("Running full suite against target B: %s", config_b.engagement.client_name)
    orchestrator_b.run()
    logger.info("Target B complete: %d findings.", len(orchestrator_b.findings))

    questionnaire_answers = _load_questionnaire(
        _resolve_nist_answers(config_a, args.config, args.nist_answers)
    )
    result = compare(
        orchestrator_a, orchestrator_b,
        target_a_name=config_a.engagement.client_name,
        target_b_name=config_b.engagement.client_name,
        questionnaire_answers=questionnaire_answers,
    )

    logger.info("Comparison complete: %d of %d techniques diverged between targets.",
                len(result.divergent_outcomes), len(result.outcomes))

    generator = ComparisonReportGenerator(auditor=config_a.engagement.auditor, language=config_a.engagement.language)
    output_path = Path(config_a.audit.output_dir) / args.compare_output
    html_path = generator.render_html(result, output_path)
    print(f"\nComparison report: {html_path}")


def main():
    parser = argparse.ArgumentParser(description="LLMRed")
    parser.add_argument("--config", required=True)
    parser.add_argument("--compare-config")
    parser.add_argument("--compare-output", default="comparison_report.html")
    parser.add_argument(
        "--nist-answers",
        help="Optional override; normally NIST evidence is auto-loaded from the main config.",
    )
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--review-findings",
        help="YAML analyst decisions keyed by evidence_id; applied before report generation.",
    )
    args = parser.parse_args()
    if args.compare_config and args.review_findings:
        parser.error("--review-findings is supported only for a single-target report.")

    try:
        config = load_config(args.config)
    except ConfigError as e:
        logger.error("Configuration error (--config): %s", e)
        sys.exit(1)

    if args.compare_config:
        try:
            config_b = load_config(args.compare_config)
        except ConfigError as e:
            logger.error("Configuration error (--compare-config): %s", e)
            sys.exit(1)
        try:
            run_comparison(config, config_b, args)
        except PolicyViolation as e:
            logger.error("Engagement policy denied execution: %s", e)
            sys.exit(3)
        except (OSError, yaml.YAMLError, ValueError) as e:
            logger.error("NIST assessment input rejected: %s", e)
            sys.exit(5)
        return

    logger.info(
        "Engagement: %s | Target: %s (%s) | Scope: %s",
        config.engagement.client_name, config.target.provider, config.target.model,
        config.engagement.scope_reference,
    )
    try:
        run_single(config, args)
    except PolicyViolation as e:
        logger.error("Engagement policy denied execution: %s", e)
        sys.exit(3)
    except ReviewFileError as e:
        logger.error("Finding review file rejected: %s", e)
        sys.exit(4)
    except (OSError, yaml.YAMLError, ValueError) as e:
        logger.error("NIST assessment input rejected: %s", e)
        sys.exit(5)


if __name__ == "__main__":
    main()