AI Pentesting Copilots

#!/usr/bin/env python3
"""AI Security Copilot — Custom implementation using OpenAI function calling."""

import json
import subprocess
import openai

client = openai.OpenAI()  # Uses OPENAI_API_KEY env var

# Define security tools the AI can invoke
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "nmap_scan",
            "description": "Run an Nmap scan against a target host or network",
            "parameters": {
                "type": "object",
                "properties": {
                    "target": {"type": "string", "description": "Target IP or hostname"},
                    "scan_type": {
                        "type": "string",
                        "enum": ["-sV", "-sC", "-sS", "-A", "-O"],
                        "description": "Nmap scan type flag"
                    },
                    "ports": {"type": "string", "description": "Port range, e.g. 1-1000"}
                },
                "required": ["target"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "nuclei_scan",
            "description": "Run Nuclei vulnerability scanner with templates",
            "parameters": {
                "type": "object",
                "properties": {
                    "target": {"type": "string", "description": "Target URL"},
                    "templates": {
                        "type": "string",
                        "description": "Template tags: cves,misconfig,exposures,takeovers"
                    },
                    "severity": {
                        "type": "string",
                        "enum": ["info", "low", "medium", "high", "critical"],
                        "description": "Minimum severity filter"
                    }
                },
                "required": ["target"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "ffuf_fuzz",
            "description": "Run ffuf directory/parameter fuzzing",
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "Target URL with FUZZ keyword"},
                    "wordlist": {"type": "string", "description": "Path to wordlist"},
                    "method": {"type": "string", "enum": ["GET", "POST"], "default": "GET"}
                },
                "required": ["url"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "analyze_response",
            "description": "Analyze an HTTP response for security issues",
            "parameters": {
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "URL to fetch and analyze"},
                    "check_headers": {"type": "boolean", "default": True},
                    "check_cookies": {"type": "boolean", "default": True}
                },
                "required": ["url"]
            }
        }
    }
]


def execute_tool(name: str, args: dict) -> str:
    """Execute a security tool and return output."""
    if name == "nmap_scan":
        cmd = ["nmap", args.get("scan_type", "-sV")]
        if "ports" in args:
            cmd += ["-p", args["ports"]]
        cmd.append(args["target"])
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        return result.stdout

    elif name == "nuclei_scan":
        cmd = ["nuclei", "-u", args["target"], "-silent"]
        if "templates" in args:
            cmd += ["-tags", args["templates"]]
        if "severity" in args:
            cmd += ["-severity", args["severity"]]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        return result.stdout

    elif name == "ffuf_fuzz":
        wordlist = args.get("wordlist", "/usr/share/seclists/Discovery/Web-Content/common.txt")
        cmd = ["ffuf", "-u", args["url"], "-w", wordlist, "-mc", "200,301,302,403"]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        return result.stdout

    elif name == "analyze_response":
        import requests
        resp = requests.get(args["url"], verify=False, timeout=10)
        analysis = f"Status: {resp.status_code}\nHeaders:\n"
        security_headers = [
            "Content-Security-Policy", "X-Frame-Options",
            "X-Content-Type-Options", "Strict-Transport-Security",
            "X-XSS-Protection", "Referrer-Policy"
        ]
        for h in security_headers:
            val = resp.headers.get(h, "MISSING")
            analysis += f"  {h}: {val}\n"
        return analysis

    return f"Unknown tool: {name}"


def run_copilot(user_message: str, target_context: str = "") -> None:
    """Run the AI security copilot with tool-calling loop."""
    system_prompt = """You are an expert penetration tester AI copilot.
You have access to security tools: nmap, nuclei, ffuf, and HTTP analysis.
Plan your approach methodically:
1. Enumerate the attack surface
2. Identify potential vulnerabilities
3. Validate findings with targeted scans
4. Provide actionable findings with severity ratings
Always explain your reasoning before executing tools."""

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Target context: {target_context}\n\n{user_message}"}
    ]

    while True:
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            tools=TOOLS,
            tool_choice="auto"
        )
        msg = response.choices[0].message
        messages.append(msg)

        if msg.tool_calls:
            for tc in msg.tool_calls:
                args = json.loads(tc.function.arguments)
                print(f"[*] Executing: {tc.function.name}({args})")
                result = execute_tool(tc.function.name, args)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": result[:8000]  # Truncate large outputs
                })
        else:
            print(msg.content)
            break


if __name__ == "__main__":
    run_copilot(
        "Perform a security assessment of the target web application.",
        target_context="Target: 10.10.10.50, Scope: *.target.htb, Auth: unauthenticated"
    )