test_chat.py

#!/usr/bin/env python3
"""
Quick test for the post-scan interactive Claude chat.
Simulates scan results so we can test the chat without running a full scan.
"""

import asyncio
import sys
from pathlib import Path

# Ensure we can import from project root
sys.path.insert(0, str(Path(__file__).parent))

from orchestrator.config_loader import load_config
from tools.claude_assistant import ClaudeAssistant


# ── Fake scan results to feed into chat ──────────
MOCK_RESULTS = [
    {
        "target": "example.com",
        "status": "complete",
        "tech_stack": [
            {"id": "nginx", "version": "1.21"},
            {"id": "react", "version": "18.2"},
            {"id": "node.js", "version": ""},
        ],
        "recon": {
            "live_hosts": ["example.com"],
            "urls_found": 42,
            "js_endpoints": ["/api/users", "/api/auth/login"],
            "forms": [{"action": "/login", "method": "POST"}],
        },
        "vulnerabilities": [
            {
                "id": "missing-x-frame-options",
                "name": "Missing X-Frame-Options Header",
                "severity": "medium",
                "url": "https://example.com",
                "description": "The X-Frame-Options header is not set, allowing clickjacking attacks.",
                "tool": "python_scanner",
            },
            {
                "id": "cors-wildcard",
                "name": "CORS Wildcard Origin",
                "severity": "high",
                "url": "https://example.com/api/users",
                "description": "Access-Control-Allow-Origin is set to * allowing any origin.",
                "tool": "python_scanner",
                "evidence": {"header": "Access-Control-Allow-Origin: *"},
            },
            {
                "id": "exposed-api-endpoint",
                "name": "Exposed API Endpoint",
                "severity": "low",
                "url": "https://example.com/api/auth/login",
                "description": "API login endpoint found via JS parsing.",
                "tool": "js_parser",
            },
        ],
        "ai_analysis": "Risk rating: MEDIUM. The CORS wildcard combined with missing clickjacking protection poses a moderate risk.",
    }
]


async def main():
    print("\n  \033[36m=== Post-Scan Chat Test ===\033[0m\n")

    # Load real config
    config = load_config(Path("config.yaml"))
    claude = ClaudeAssistant(config)

    if not claude.available:
        print("  \033[31m[!] Claude assistant is not available. Check config.yaml.\033[0m")
        return

    print("  \033[32m[OK] Claude assistant initialized.\033[0m")

    # Build scan context
    scan_context = {"results": MOCK_RESULTS}
    all_findings = MOCK_RESULTS[0]["vulnerabilities"]
    all_tech = MOCK_RESULTS[0]["tech_stack"]

    # Show summary
    print(f"  \033[90mLoaded {len(all_findings)} mock findings.\033[0m")
    for i, f in enumerate(all_findings, 1):
        print(f"  \033[90m  {i}. [{f['severity'].upper()}] {f['id']}\033[0m")
    print()

    # Reset and start chat
    claude.reset_chat()

    print("  \033[36mTesting chat with a question...\033[0m\n")

    reply = await claude.chat(
        "What are the most critical findings and what should I fix first?",
        scan_context=scan_context,
    )

    if "[AUTH_REQUIRED]" in reply:
        print(f"  \033[33m{reply}\033[0m")
        input("  \033[33m>\033[0m Press Enter after authenticating in your browser... ")
        reply = await claude.chat(
            "What are the most critical findings and what should I fix first?",
            scan_context=scan_context,
        )

    if reply and "[AUTH_REQUIRED]" not in reply:
        # Indent for display
        for line in reply.splitlines():
            print(f"  {line}")
        print()
        print("  \033[32m[OK] Chat is working!\033[0m")
    else:
        print("  \033[31m[!] No response or auth still failing.\033[0m")
        return

    # Test remediation
    print("\n  \033[36mTesting /remediate for finding #2 (CORS wildcard)...\033[0m\n")

    remediation = await claude.get_remediation(all_findings[1], 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(all_findings[1], all_tech)

    if remediation:
        for line in remediation.splitlines():
            print(f"  {line}")
        print()
        print("  \033[32m[OK] Remediation advice working!\033[0m")
    else:
        print("  \033[31m[!] No remediation response.\033[0m")

    print("\n  \033[32m=== All tests passed! ===\033[0m\n")


if __name__ == "__main__":
    asyncio.run(main())