Using AI for Vulnerability Research

from dataclasses import dataclass
from enum import Enum
 
class AIEffectiveness(Enum):
    HIGH = "high"       # AI significantly accelerates this task
    MEDIUM = "medium"   # AI helps but requires heavy guidance
    LOW = "low"         # AI provides marginal value
    NEGATIVE = "negative"  # AI actively misleads
 
@dataclass
class ResearchPhase:
    phase: str
    description: str
    ai_effectiveness: AIEffectiveness
    best_tools: list[str]
    ai_limitations: list[str]
 
RESEARCH_PHASES = [
    ResearchPhase(
        phase="Attack Surface Mapping",
        description="Identifying entry points, data flows, and trust boundaries",
        ai_effectiveness=AIEffectiveness.HIGH,
        best_tools=["Claude Code", "Cursor", "Aider"],
        ai_limitations=[
            "May miss non-obvious entry points (IPC, shared memory)",
            "Cannot analyze running processes or dynamic behavior",
        ],
    ),
    ResearchPhase(
        phase="Code Comprehension",
        description="Understanding unfamiliar codebases, protocols, or file formats",
        ai_effectiveness=AIEffectiveness.HIGH,
        best_tools=["Claude Code", "Cursor"],
        ai_limitations=[
            "May hallucinate API behavior not in context window",
            "Cannot understand undocumented proprietary protocols",
        ],
    ),
    ResearchPhase(
        phase="Vulnerability Pattern Recognition",
        description="Identifying potential vulnerability classes in code",
        ai_effectiveness=AIEffectiveness.MEDIUM,
        best_tools=["Semgrep", "CodeQL", "Claude Code"],
        ai_limitations=[
            "High false positive rate for complex vulnerability classes",
            "May miss subtle logic bugs that require deep understanding",
            "Often misidentifies severity of findings",
        ],
    ),
    ResearchPhase(
        phase="Fuzzer Development",
        description="Writing targeted fuzzers for specific attack surfaces",
        ai_effectiveness=AIEffectiveness.HIGH,
        best_tools=["Claude Code", "GitHub Copilot", "Cursor"],
        ai_limitations=[
            "Generated fuzzers may not cover important code paths",
            "May not handle edge cases in protocol/format parsing",
        ],
    ),
    ResearchPhase(
        phase="Crash Triage",
        description="Analyzing crashes, determining root cause and exploitability",
        ai_effectiveness=AIEffectiveness.MEDIUM,
        best_tools=["Claude Code"],
        ai_limitations=[
            "Cannot run debugger or analyze memory state",
            "May misidentify root cause if crash is complex",
            "Limited ability to assess exploitability without runtime info",
        ],
    ),
    ResearchPhase(
        phase="Exploit Development",
        description="Building proof-of-concept exploits",
        ai_effectiveness=AIEffectiveness.MEDIUM,
        best_tools=["Claude Code", "Cursor"],
        ai_limitations=[
            "Exploit code often subtly incorrect (wrong offsets, encoding)",
            "Limited understanding of memory layout and ASLR bypass",
            "Cannot test exploit code without execution environment",
        ],
    ),
    ResearchPhase(
        phase="Vulnerability Reporting",
        description="Writing vulnerability reports and advisories",
        ai_effectiveness=AIEffectiveness.HIGH,
        best_tools=["Claude Code", "any LLM"],
        ai_limitations=[
            "May overstate or understate severity",
            "CVSS scoring requires human validation",
        ],
    ),
]

# AI-assisted code audit workflow
 
import subprocess
from pathlib import Path
 
class AIAssistedCodeAudit:
    """Framework for AI-assisted source code security audit."""
 
    def __init__(self, project_path: str):
        self.project_path = Path(project_path)
 
    def identify_high_risk_areas(self) -> dict:
        """Use heuristics to identify code areas that need AI-assisted review."""
        high_risk_patterns = {
            "input_handlers": {
                "patterns": [
                    r"request\.(get|post|json|form|args|data)",
                    r"sys\.argv",
                    r"input\(",
                    r"socket\.recv",
                    r"read\(",
                ],
                "risk": "External input entry points",
            },
            "crypto_operations": {
                "patterns": [
                    r"(hashlib|hmac|cryptography|Crypto|ssl)\.",
                    r"(encrypt|decrypt|sign|verify|hash)",
                    r"(AES|RSA|DES|MD5|SHA)",
                ],
                "risk": "Cryptographic operations may have implementation flaws",
            },
            "auth_boundaries": {
                "patterns": [
                    r"(login|authenticate|authorize|session|token|jwt)",
                    r"@(login_required|auth|permission)",
                    r"(password|credential|secret)",
                ],
                "risk": "Authentication and authorization logic",
            },
            "command_execution": {
                "patterns": [
                    r"(subprocess|os\.system|os\.popen|exec|eval)",
                    r"(shell=True|Popen)",
                    r"(exec\(|eval\(|compile\()",
                ],
                "risk": "Command injection and code execution",
            },
            "deserialization": {
                "patterns": [
                    r"(pickle\.load|yaml\.load|json\.loads)",
                    r"(deserialize|unmarshal|unpack)",
                    r"(fromstring|parse.*xml)",
                ],
                "risk": "Deserialization of untrusted data",
            },
            "file_operations": {
                "patterns": [
                    r"(open\(.*\+|os\.path\.join|shutil)",
                    r"(upload|download|file_path|filename)",
                    r"(mkdir|rmdir|unlink|chmod)",
                ],
                "risk": "Path traversal and file operation vulnerabilities",
            },
        }
 
        results = {}
        for category, config in high_risk_patterns.items():
            matches = []
            for pattern in config["patterns"]:
                try:
                    result = subprocess.run(
                        ["grep", "-rnI", "-E", pattern, str(self.project_path),
                         "--include=*.py", "--include=*.js", "--include=*.ts",
                         "--exclude-dir=node_modules", "--exclude-dir=.git",
                         "--exclude-dir=venv"],
                        capture_output=True, text=True,
                    )
                    for line in result.stdout.strip().split("\n"):
                        if line:
                            matches.append(line)
                except Exception:
                    pass
 
            results[category] = {
                "risk_description": config["risk"],
                "match_count": len(matches),
                "sample_matches": matches[:5],
            }
 
        return results
 
    def generate_audit_prompts(self, high_risk_areas: dict) -> list[dict]:
        """Generate targeted prompts for AI-assisted review of each area."""
        prompts = []
 
        prompt_templates = {
            "input_handlers": (
                "Review this code for input validation vulnerabilities. "
                "Check for: SQL injection, command injection, XSS, path traversal, "
                "SSRF, and deserialization attacks. For each input handler, trace "
                "the data flow from entry to processing and identify missing "
                "sanitization or validation."
            ),
            "crypto_operations": (
                "Review this cryptographic code for implementation flaws. "
                "Check for: weak algorithms (MD5, SHA1 for security), "
                "static IVs/nonces, missing integrity checks, improper "
                "key derivation, timing side channels, and insecure random "
                "number generation."
            ),
            "auth_boundaries": (
                "Review this authentication and authorization code. "
                "Check for: authentication bypass, broken access control, "
                "session fixation, token leakage, password storage flaws, "
                "missing authorization checks on sensitive operations, "
                "and IDOR vulnerabilities."
            ),
            "command_execution": (
                "Review this code for command injection vulnerabilities. "
                "Check for: unsanitized input in subprocess calls, "
                "shell=True with user input, eval/exec with user input, "
                "and environment variable injection."
            ),
            "deserialization": (
                "Review this code for insecure deserialization. "
                "Check for: pickle.load on untrusted data, yaml.load "
                "without SafeLoader, XML external entity (XXE) attacks, "
                "and prototype pollution."
            ),
            "file_operations": (
                "Review this code for file operation vulnerabilities. "
                "Check for: path traversal (../), symlink attacks, "
                "TOCTOU race conditions, insecure temporary files, "
                "and unrestricted file upload."
            ),
        }
 
        for category, area in high_risk_areas.items():
            if area["match_count"] > 0:
                prompts.append({
                    "category": category,
                    "prompt": prompt_templates.get(category, "Review this code for security issues."),
                    "files_to_review": [
                        m.split(":")[0] for m in area["sample_matches"]
                    ],
                    "match_count": area["match_count"],
                })
 
        return prompts

# Example: Using AI to generate a targeted fuzzer for a parser
 
# Step 1: Provide the AI with the target function signature and context
TARGET_CONTEXT = """
The target is a custom binary protocol parser in Python:
 
def parse_message(data: bytes) -> dict:
    if len(data) < 8:
        raise ValueError("Message too short")
    magic = struct.unpack('>H', data[0:2])[0]
    if magic != 0xCAFE:
        raise ValueError("Invalid magic")
    version = data[2]
    msg_type = data[3]
    payload_len = struct.unpack('>I', data[4:8])[0]
    payload = data[8:8+payload_len]
    # ... further parsing based on msg_type
"""
 
# Step 2: AI generates a structured fuzzer
# This is the type of output AI can produce effectively:
 
import atheris
import sys
import struct
 
# AI-generated fuzzer for the binary protocol parser
@atheris.instrument_func
def fuzz_parse_message(data):
    """Fuzz the message parser with structure-aware inputs."""
    fdp = atheris.FuzzedDataProvider(data)
 
    # Generate valid-ish messages that pass initial checks
    # This structure-awareness is where AI adds value
    magic = fdp.ConsumeIntInRange(0, 0xFFFF)
    version = fdp.ConsumeIntInRange(0, 255)
    msg_type = fdp.ConsumeIntInRange(0, 255)
    payload = fdp.ConsumeBytes(fdp.ConsumeIntInRange(0, 1024))
 
    # Construct message with proper structure
    message = struct.pack('>HBB', magic, version, msg_type)
    message += struct.pack('>I', len(payload))
    message += payload
 
    try:
        # Import the target module
        from target import parse_message
        parse_message(message)
    except (ValueError, struct.error, IndexError, KeyError):
        # Expected exceptions from invalid input
        pass
    except Exception as e:
        # Unexpected exceptions may indicate bugs
        raise
 
def main():
    atheris.Setup(sys.argv, fuzz_parse_message)
    atheris.Fuzz()
 
if __name__ == "__main__":
    main()

# Templates for AI-assisted fuzzer generation across target types
 
FUZZER_TEMPLATES = {
    "http_api_fuzzer": {
        "description": "Fuzz HTTP API endpoints with mutated inputs",
        "ai_prompt": (
            "Generate a Python fuzzer using the requests library that tests "
            "this API endpoint. The fuzzer should: "
            "1. Start with valid request templates "
            "2. Mutate headers, parameters, body fields, and content types "
            "3. Check for 500 errors, timing anomalies, and unexpected responses "
            "4. Log all interesting responses for manual review"
        ),
        "example_output": '''
import requests
import random
import string
import json
import time
from typing import Optional
 
class APIFuzzer:
    """AI-generated API fuzzer with structured mutations."""
 
    def __init__(self, base_url: str, auth_token: Optional[str] = None):
        self.base_url = base_url
        self.session = requests.Session()
        if auth_token:
            self.session.headers["Authorization"] = f"Bearer {auth_token}"
        self.interesting_responses: list[dict] = []
 
    def _mutate_string(self, value: str) -> str:
        """Apply random string mutations."""
        mutations = [
            lambda v: v + "'" + "OR 1=1--",  # SQL injection
            lambda v: "<script>alert(1)</script>",  # XSS
            lambda v: "../" * 10 + "etc/passwd",  # Path traversal
            lambda v: "A" * 10000,  # Buffer overflow
            lambda v: v + "\\x00" + "injected",  # Null byte
            lambda v: "{{7*7}}",  # Template injection
            lambda v: "${7*7}",  # Expression injection
            lambda v: "",  # Empty string
            lambda v: " ",  # Whitespace only
        ]
        return random.choice(mutations)(value)
 
    def _mutate_json(self, data: dict) -> dict:
        """Apply mutations to JSON body fields."""
        mutated = data.copy()
        if not mutated:
            return mutated
 
        key = random.choice(list(mutated.keys()))
        value = mutated[key]
 
        if isinstance(value, str):
            mutated[key] = self._mutate_string(value)
        elif isinstance(value, int):
            mutated[key] = random.choice([0, -1, 2**31, 2**63, None])
        elif isinstance(value, list):
            mutated[key] = [self._mutate_string("x")] * 1000
        elif isinstance(value, bool):
            mutated[key] = not value
 
        return mutated
 
    def fuzz_endpoint(
        self,
        path: str,
        method: str,
        base_body: Optional[dict] = None,
        iterations: int = 100,
    ) -> list[dict]:
        """Fuzz a specific endpoint with mutations."""
        url = f"{self.base_url}{path}"
 
        for i in range(iterations):
            try:
                body = self._mutate_json(base_body) if base_body else None
                start = time.time()
 
                if method == "GET":
                    response = self.session.get(url, timeout=10)
                elif method == "POST":
                    response = self.session.post(url, json=body, timeout=10)
                elif method == "PUT":
                    response = self.session.put(url, json=body, timeout=10)
                else:
                    continue
 
                elapsed = time.time() - start
 
                # Flag interesting responses
                if (response.status_code >= 500
                    or elapsed > 5.0
                    or "error" in response.text.lower()
                    and "traceback" in response.text.lower()):
 
                    self.interesting_responses.append({
                        "iteration": i,
                        "url": url,
                        "method": method,
                        "body": body,
                        "status": response.status_code,
                        "elapsed": elapsed,
                        "response_preview": response.text[:500],
                    })
 
            except requests.exceptions.Timeout:
                self.interesting_responses.append({
                    "iteration": i,
                    "url": url,
                    "method": method,
                    "body": body,
                    "issue": "timeout",
                })
            except requests.exceptions.RequestException:
                pass
 
        return self.interesting_responses
''',
    },
 
    "file_format_fuzzer": {
        "description": "Fuzz file format parsers",
        "ai_prompt": (
            "Generate a fuzzer that mutates a valid file and feeds "
            "mutations to a parser, checking for crashes."
        ),
    },
}

# Using AI to triage and analyze crashes found by fuzzing
 
class CrashTriager:
    """AI-assisted crash triage and analysis."""
 
    def __init__(self):
        self.triaged_crashes: list[dict] = []
 
    def parse_crash_report(self, crash_output: str) -> dict:
        """Parse crash information from various sources."""
        import re
 
        crash_info = {
            "raw_output": crash_output,
            "exception_type": None,
            "exception_message": None,
            "stack_frames": [],
            "source_files": set(),
        }
 
        # Extract Python traceback
        tb_match = re.search(
            r"Traceback \(most recent call last\):\n(.*?)(\w+Error.*?)$",
            crash_output, re.DOTALL | re.MULTILINE,
        )
        if tb_match:
            frames_text = tb_match.group(1)
            crash_info["exception_type"] = tb_match.group(2).split(":")[0].strip()
            crash_info["exception_message"] = tb_match.group(2).strip()
 
            # Parse individual frames
            frame_pattern = r'File "([^"]+)", line (\d+), in (\w+)'
            for match in re.finditer(frame_pattern, frames_text):
                crash_info["stack_frames"].append({
                    "file": match.group(1),
                    "line": int(match.group(2)),
                    "function": match.group(3),
                })
                crash_info["source_files"].add(match.group(1))
 
        crash_info["source_files"] = list(crash_info["source_files"])
        return crash_info
 
    def generate_ai_triage_prompt(self, crash_info: dict) -> str:
        """Generate a prompt for AI-assisted crash triage."""
        prompt = f"""Analyze this crash for security implications:
 
Exception: {crash_info.get('exception_type')}: {crash_info.get('exception_message')}
 
Stack trace:
"""
        for frame in crash_info.get("stack_frames", []):
            prompt += f"  {frame['file']}:{frame['line']} in {frame['function']}\n"
 
        prompt += """
Questions to answer:
1. What is the root cause of this crash?
2. Is this crash security-relevant? Could it lead to:
   - Code execution?
   - Information disclosure?
   - Denial of service?
   - Authentication bypass?
3. What CWE does this most closely map to?
4. What would a proof-of-concept exploit look like?
5. How should this be fixed?
"""
        return prompt
 
    def assess_exploitability(self, crash_info: dict) -> dict:
        """Heuristic exploitability assessment."""
        score = 0
        factors = []
 
        exception_type = crash_info.get("exception_type", "")
 
        # Higher exploitability indicators
        if "BufferError" in exception_type or "OverflowError" in exception_type:
            score += 3
            factors.append("Memory corruption indicator")
 
        if "PermissionError" in exception_type or "AuthError" in exception_type:
            score += 2
            factors.append("Permission/auth boundary crossed")
 
        if any("deserializ" in f.get("function", "").lower()
               for f in crash_info.get("stack_frames", [])):
            score += 3
            factors.append("Crash in deserialization code")
 
        if any("parse" in f.get("function", "").lower()
               for f in crash_info.get("stack_frames", [])):
            score += 1
            factors.append("Crash in parser code")
 
        exploitability = "low"
        if score >= 3:
            exploitability = "high"
        elif score >= 2:
            exploitability = "medium"
 
        return {
            "exploitability": exploitability,
            "score": score,
            "factors": factors,
        }

RESPONSIBLE_AI_VULN_RESEARCH = {
    "authorized_research": {
        "description": "Research conducted with proper authorization",
        "requirements": [
            "Written authorization from system owner (bug bounty, engagement letter)",
            "Defined scope of testing",
            "Responsible disclosure timeline agreed upon",
            "No access to production user data",
        ],
        "ai_appropriate": True,
    },
    "coordinated_disclosure": {
        "description": "Working with vendors to fix vulnerabilities",
        "requirements": [
            "Report to vendor before public disclosure",
            "Follow vendor's preferred disclosure process",
            "Allow reasonable time for fix development (90 days standard)",
            "Provide proof-of-concept, not weaponized exploit",
        ],
        "ai_appropriate": True,
    },
    "weaponization": {
        "description": "Creating fully weaponized exploits for unauthorized use",
        "requirements": [],
        "ai_appropriate": False,
        "note": (
            "AI tools should not be used to create weaponized exploits "
            "for unauthorized access. This violates laws including CFAA "
            "and most AI tool acceptable use policies."
        ),
    },
}
 
# What AI tools CAN and CANNOT do for vuln research
AI_CAPABILITIES_HONEST = {
    "genuinely_helpful": [
        "Explaining unfamiliar code patterns and protocols",
        "Generating boilerplate for fuzzers and test harnesses",
        "Writing Semgrep/CodeQL rules for specific vulnerability patterns",
        "Drafting vulnerability reports and advisories",
        "Suggesting related CWEs and attack techniques",
        "Converting between programming languages for PoC porting",
    ],
    "helpful_with_guidance": [
        "Identifying potential vulnerability classes in source code",
        "Generating proof-of-concept exploit code (requires manual tuning)",
        "Analyzing crash dumps and stack traces",
        "Suggesting attack vectors for a given attack surface",
    ],
    "unreliable": [
        "Finding novel zero-day vulnerabilities without guidance",
        "Accurately assessing exploitability of memory corruption bugs",
        "Generating working exploits for complex vulnerability chains",
        "Bypassing modern mitigations (ASLR, CFI, sandboxes)",
        "Understanding proprietary/undocumented protocols",
    ],
}

# Complete AI-augmented vulnerability research workflow
 
RESEARCH_WORKFLOW = {
    "step_1_reconnaissance": {
        "human_task": "Define research target and scope",
        "ai_task": "Map attack surface, identify entry points, list technologies",
        "tools": ["Claude Code for codebase analysis", "Semgrep for pattern matching"],
        "output": "Attack surface map with prioritized areas",
    },
    "step_2_analysis": {
        "human_task": "Guide AI to focus on high-risk areas",
        "ai_task": "Deep code review of identified areas",
        "tools": ["CodeQL for dataflow analysis", "AI code review prompts"],
        "output": "List of potential vulnerabilities with confidence levels",
    },
    "step_3_validation": {
        "human_task": "Verify AI findings, filter false positives",
        "ai_task": "Generate test cases and PoC code for each finding",
        "tools": ["AI-generated fuzzers", "manual debugging"],
        "output": "Confirmed vulnerabilities with PoC",
    },
    "step_4_exploitation": {
        "human_task": "Assess exploitability, develop reliable PoC",
        "ai_task": "Generate exploit scaffolding, convert between formats",
        "tools": ["AI for boilerplate", "manual exploit refinement"],
        "output": "Working PoC demonstrating impact",
    },
    "step_5_reporting": {
        "human_task": "Validate severity, determine CVSS score",
        "ai_task": "Draft vulnerability report, suggest mitigations",
        "tools": ["AI for report writing", "CVSS calculator"],
        "output": "Vulnerability report ready for submission",
    },
}