aiptx 2.0.7__py3-none-any.whl

aipt_v2/agents/exploit_agent.py

@@ -0,0 +1,688 @@
+"""
+AIPT Exploitation Reasoning Agent
+
+An LLM-powered agent that reasons through exploitation step-by-step:
+- Analyzes vulnerability context
+- Plans exploitation strategy
+- Executes actions and learns from results
+- Adapts approach based on feedback
+
+This provides intelligent, adaptive exploitation beyond script-based attacks.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Callable, Optional
+
+from aipt_v2.models.findings import Finding, Severity, VulnerabilityType
+
+logger = logging.getLogger(__name__)
+
+
+EXPLOIT_REASONING_PROMPT = """You are an expert penetration tester exploiting a vulnerability.
+
+## Vulnerability Details
+- **Title**: {title}
+- **Type**: {vuln_type}
+- **Severity**: {severity}
+- **URL**: {url}
+- **Parameter**: {parameter}
+- **Evidence**: {evidence}
+
+## Environment Context
+- **Technology Stack**: {tech_stack}
+- **WAF Detected**: {waf}
+- **Authentication**: {auth_status}
+
+## Exploitation History
+{history}
+
+## Your Objective
+{objective}
+
+## Available Actions
+You can perform ONE of these actions:
+
+1. **send_request**: Send an HTTP request
+   - method: GET|POST|PUT|DELETE
+   - url: target URL
+   - headers: dict of headers
+   - body: request body
+   - params: query parameters
+
+2. **test_payload**: Test a specific payload
+   - payload: the payload string
+   - encoding: none|url|double_url|base64|unicode
+
+3. **analyze_response**: Analyze the last response
+   - check_for: what to look for (error, success_indicator, data)
+
+4. **extract_data**: Extract data from response
+   - pattern: regex pattern to extract
+   - data_type: credentials|tokens|database|files
+
+5. **escalate**: Attempt privilege escalation
+   - technique: the escalation technique
+
+6. **conclude**: End exploitation with result
+   - success: true|false
+   - evidence: proof of exploitation
+   - impact: description of achieved impact
+
+## Think Step by Step
+1. What is the current state of the exploitation?
+2. What obstacles exist (WAF, auth, etc.)?
+3. What should you try next?
+4. What do you expect to happen?
+
+## Response Format (JSON)
+```json
+{{
+    "reasoning": "Your step-by-step reasoning",
+    "action": "action_name",
+    "parameters": {{
+        "param1": "value1"
+    }},
+    "expected_outcome": "What you expect to happen",
+    "fallback_plan": "What to try if this fails"
+}}
+```"""
+
+
+class ExploitAction(Enum):
+    """Actions the exploit agent can take."""
+    SEND_REQUEST = "send_request"
+    TEST_PAYLOAD = "test_payload"
+    ANALYZE_RESPONSE = "analyze_response"
+    EXTRACT_DATA = "extract_data"
+    ESCALATE = "escalate"
+    CONCLUDE = "conclude"
+
+
+@dataclass
+class ActionResult:
+    """Result of an exploitation action."""
+    action: ExploitAction
+    success: bool
+    response_code: Optional[int] = None
+    response_body: Optional[str] = None
+    response_time_ms: Optional[int] = None
+    extracted_data: Optional[dict] = None
+    error: Optional[str] = None
+    timestamp: datetime = field(default_factory=datetime.utcnow)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "action": self.action.value,
+            "success": self.success,
+            "response_code": self.response_code,
+            "response_body": self.response_body[:500] if self.response_body else None,
+            "response_time_ms": self.response_time_ms,
+            "extracted_data": self.extracted_data,
+            "error": self.error,
+        }
+
+
+@dataclass
+class ExploitStep:
+    """A single step in the exploitation process."""
+    step_number: int
+    reasoning: str
+    action: ExploitAction
+    parameters: dict[str, Any]
+    expected_outcome: str
+    result: Optional[ActionResult] = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "step": self.step_number,
+            "reasoning": self.reasoning,
+            "action": self.action.value,
+            "parameters": self.parameters,
+            "expected_outcome": self.expected_outcome,
+            "result": self.result.to_dict() if self.result else None,
+        }
+
+
+@dataclass
+class ExploitResult:
+    """Final result of exploitation attempt."""
+    success: bool
+    finding: Finding
+    steps: list[ExploitStep]
+    evidence: str
+    impact_achieved: str
+    exploitation_time_seconds: float
+    total_requests: int
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "success": self.success,
+            "finding_title": self.finding.title,
+            "finding_url": self.finding.url,
+            "steps": [s.to_dict() for s in self.steps],
+            "evidence": self.evidence,
+            "impact_achieved": self.impact_achieved,
+            "exploitation_time_seconds": self.exploitation_time_seconds,
+            "total_requests": self.total_requests,
+        }
+
+
+class ExploitReasoningAgent:
+    """
+    LLM-powered agent that reasons through exploitation.
+
+    Uses a Think-Act-Observe loop to intelligently exploit vulnerabilities:
+    1. THINK: Analyze current state and plan next action
+    2. ACT: Execute the planned action
+    3. OBSERVE: Analyze the result and update state
+    4. REPEAT until success or max attempts
+
+    Example:
+        agent = ExploitReasoningAgent()
+
+        result = await agent.exploit(
+            finding=sqli_finding,
+            objective="Extract database credentials",
+            max_attempts=15
+        )
+
+        if result.success:
+            print(f"Exploitation successful!")
+            print(f"Evidence: {result.evidence}")
+    """
+
+    def __init__(
+        self,
+        llm_provider: str = "anthropic",
+        llm_model: str = "claude-3-5-sonnet-20241022",
+        http_client: Any = None,
+    ):
+        self.llm_provider = llm_provider
+        self.llm_model = llm_model
+        self.http_client = http_client
+        self._llm = None
+        self._action_handlers: dict[ExploitAction, Callable] = {}
+        self._setup_default_handlers()
+
+    async def _get_llm(self):
+        """Get or create LLM client."""
+        if self._llm is None:
+            try:
+                import litellm
+                self._llm = litellm
+            except ImportError:
+                logger.error("litellm not installed")
+                return None
+        return self._llm
+
+    def _setup_default_handlers(self):
+        """Set up default action handlers."""
+        self._action_handlers = {
+            ExploitAction.SEND_REQUEST: self._handle_send_request,
+            ExploitAction.TEST_PAYLOAD: self._handle_test_payload,
+            ExploitAction.ANALYZE_RESPONSE: self._handle_analyze_response,
+            ExploitAction.EXTRACT_DATA: self._handle_extract_data,
+            ExploitAction.ESCALATE: self._handle_escalate,
+            ExploitAction.CONCLUDE: self._handle_conclude,
+        }
+
+    def register_handler(self, action: ExploitAction, handler: Callable):
+        """Register a custom handler for an action."""
+        self._action_handlers[action] = handler
+
+    async def exploit(
+        self,
+        finding: Finding,
+        objective: str = "Achieve successful exploitation",
+        tech_stack: str = None,
+        waf: str = None,
+        auth_status: str = "None",
+        max_attempts: int = 15,
+    ) -> ExploitResult:
+        """
+        Attempt to exploit a vulnerability using reasoning.
+
+        Args:
+            finding: The vulnerability finding to exploit
+            objective: The exploitation objective
+            tech_stack: Target technology stack
+            waf: Detected WAF name
+            auth_status: Current authentication status
+            max_attempts: Maximum exploitation attempts
+
+        Returns:
+            ExploitResult with exploitation outcome
+        """
+        start_time = datetime.utcnow()
+        steps: list[ExploitStep] = []
+        total_requests = 0
+        last_response = None
+
+        llm = await self._get_llm()
+        if llm is None:
+            return self._failed_result(finding, "LLM not available", steps, 0, 0)
+
+        # Build exploitation context
+        context = {
+            "title": finding.title,
+            "vuln_type": finding.vuln_type.value,
+            "severity": finding.severity.value,
+            "url": finding.url,
+            "parameter": finding.parameter or "N/A",
+            "evidence": finding.evidence[:500] if finding.evidence else "None",
+            "tech_stack": tech_stack or "Unknown",
+            "waf": waf or "None detected",
+            "auth_status": auth_status,
+            "objective": objective,
+        }
+
+        logger.info(f"Starting exploitation of: {finding.title}")
+
+        for attempt in range(max_attempts):
+            # Build history from previous steps
+            history = self._format_history(steps)
+            context["history"] = history
+
+            # Get next action from LLM
+            try:
+                prompt = EXPLOIT_REASONING_PROMPT.format(**context)
+                response = await self._call_llm(prompt)
+                action_data = self._parse_action(response)
+            except Exception as e:
+                logger.warning(f"LLM call failed at step {attempt + 1}: {e}")
+                continue
+
+            # Create step
+            try:
+                action = ExploitAction(action_data.get("action", "conclude"))
+            except ValueError:
+                action = ExploitAction.CONCLUDE
+
+            step = ExploitStep(
+                step_number=attempt + 1,
+                reasoning=action_data.get("reasoning", ""),
+                action=action,
+                parameters=action_data.get("parameters", {}),
+                expected_outcome=action_data.get("expected_outcome", ""),
+            )
+
+            # Execute action
+            handler = self._action_handlers.get(action)
+            if handler:
+                try:
+                    result = await handler(
+                        finding=finding,
+                        parameters=step.parameters,
+                        last_response=last_response,
+                    )
+                    step.result = result
+                    if result.response_code:
+                        total_requests += 1
+                    last_response = result
+                except Exception as e:
+                    step.result = ActionResult(
+                        action=action,
+                        success=False,
+                        error=str(e),
+                    )
+
+            steps.append(step)
+
+            # Check for conclusion
+            if action == ExploitAction.CONCLUDE:
+                success = step.parameters.get("success", False)
+                evidence = step.parameters.get("evidence", "")
+                impact = step.parameters.get("impact", "")
+
+                elapsed = (datetime.utcnow() - start_time).total_seconds()
+
+                return ExploitResult(
+                    success=success,
+                    finding=finding,
+                    steps=steps,
+                    evidence=evidence,
+                    impact_achieved=impact,
+                    exploitation_time_seconds=elapsed,
+                    total_requests=total_requests,
+                )
+
+            # Log progress
+            logger.debug(f"Step {attempt + 1}: {action.value} - "
+                        f"{'Success' if step.result and step.result.success else 'Failed'}")
+
+        # Max attempts reached
+        elapsed = (datetime.utcnow() - start_time).total_seconds()
+        return ExploitResult(
+            success=False,
+            finding=finding,
+            steps=steps,
+            evidence="Max attempts reached without successful exploitation",
+            impact_achieved="None - exploitation unsuccessful",
+            exploitation_time_seconds=elapsed,
+            total_requests=total_requests,
+        )
+
+    async def _call_llm(self, prompt: str) -> str:
+        """Call LLM for reasoning."""
+        llm = await self._get_llm()
+
+        model_str = f"{self.llm_provider}/{self.llm_model}"
+        if self.llm_provider == "anthropic" and not self.llm_model.startswith("anthropic/"):
+            model_str = f"anthropic/{self.llm_model}"
+        elif self.llm_provider == "openai" and not self.llm_model.startswith("openai/"):
+            model_str = f"openai/{self.llm_model}"
+
+        response = await llm.acompletion(
+            model=model_str,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=1500,
+            temperature=0.3,
+        )
+
+        return response.choices[0].message.content
+
+    def _parse_action(self, response: str) -> dict[str, Any]:
+        """Parse action from LLM response."""
+        try:
+            # Extract JSON from response
+            json_start = response.find("{")
+            json_end = response.rfind("}") + 1
+            if json_start >= 0 and json_end > json_start:
+                return json.loads(response[json_start:json_end])
+        except (json.JSONDecodeError, ValueError):
+            pass
+
+        # Default to conclude if parsing fails
+        return {
+            "reasoning": "Failed to parse response",
+            "action": "conclude",
+            "parameters": {"success": False, "evidence": "Parse error"},
+        }
+
+    def _format_history(self, steps: list[ExploitStep]) -> str:
+        """Format exploitation history for LLM context."""
+        if not steps:
+            return "No previous actions taken."
+
+        lines = ["Previous exploitation attempts:"]
+        for step in steps[-5:]:  # Only last 5 steps for context
+            result_str = "Success" if step.result and step.result.success else "Failed"
+            lines.append(f"- Step {step.step_number}: {step.action.value} - {result_str}")
+            if step.result and step.result.error:
+                lines.append(f"  Error: {step.result.error}")
+            if step.result and step.result.response_code:
+                lines.append(f"  Response: HTTP {step.result.response_code}")
+
+        return "\n".join(lines)
+
+    # ==================== Action Handlers ====================
+
+    async def _handle_send_request(
+        self,
+        finding: Finding,
+        parameters: dict[str, Any],
+        last_response: Optional[ActionResult],
+    ) -> ActionResult:
+        """Handle send_request action."""
+        import time
+
+        method = parameters.get("method", "GET").upper()
+        url = parameters.get("url", finding.url)
+        headers = parameters.get("headers", {})
+        body = parameters.get("body")
+        params = parameters.get("params", {})
+
+        try:
+            import httpx
+
+            start = time.time()
+            async with httpx.AsyncClient(verify=False, timeout=30) as client:
+                response = await client.request(
+                    method=method,
+                    url=url,
+                    headers=headers,
+                    content=body,
+                    params=params,
+                )
+            elapsed_ms = int((time.time() - start) * 1000)
+
+            return ActionResult(
+                action=ExploitAction.SEND_REQUEST,
+                success=response.status_code < 400,
+                response_code=response.status_code,
+                response_body=response.text[:2000],
+                response_time_ms=elapsed_ms,
+            )
+
+        except ImportError:
+            return ActionResult(
+                action=ExploitAction.SEND_REQUEST,
+                success=False,
+                error="httpx not installed",
+            )
+        except Exception as e:
+            return ActionResult(
+                action=ExploitAction.SEND_REQUEST,
+                success=False,
+                error=str(e),
+            )
+
+    async def _handle_test_payload(
+        self,
+        finding: Finding,
+        parameters: dict[str, Any],
+        last_response: Optional[ActionResult],
+    ) -> ActionResult:
+        """Handle test_payload action."""
+        import time
+        import urllib.parse
+
+        payload = parameters.get("payload", "")
+        encoding = parameters.get("encoding", "none")
+
+        # Apply encoding
+        if encoding == "url":
+            payload = urllib.parse.quote(payload)
+        elif encoding == "double_url":
+            payload = urllib.parse.quote(urllib.parse.quote(payload))
+        elif encoding == "base64":
+            import base64
+            payload = base64.b64encode(payload.encode()).decode()
+
+        # Build request with payload
+        url = finding.url
+        if finding.parameter:
+            if "?" in url:
+                url = f"{url}&{finding.parameter}={payload}"
+            else:
+                url = f"{url}?{finding.parameter}={payload}"
+
+        try:
+            import httpx
+
+            start = time.time()
+            async with httpx.AsyncClient(verify=False, timeout=30) as client:
+                response = await client.get(url)
+            elapsed_ms = int((time.time() - start) * 1000)
+
+            # Check for success indicators
+            success_indicators = [
+                response.status_code == 200,
+                "error" in response.text.lower(),
+                "sql" in response.text.lower(),
+                "syntax" in response.text.lower(),
+            ]
+
+            return ActionResult(
+                action=ExploitAction.TEST_PAYLOAD,
+                success=any(success_indicators),
+                response_code=response.status_code,
+                response_body=response.text[:2000],
+                response_time_ms=elapsed_ms,
+            )
+
+        except Exception as e:
+            return ActionResult(
+                action=ExploitAction.TEST_PAYLOAD,
+                success=False,
+                error=str(e),
+            )
+
+    async def _handle_analyze_response(
+        self,
+        finding: Finding,
+        parameters: dict[str, Any],
+        last_response: Optional[ActionResult],
+    ) -> ActionResult:
+        """Handle analyze_response action."""
+        check_for = parameters.get("check_for", "error")
+
+        if not last_response or not last_response.response_body:
+            return ActionResult(
+                action=ExploitAction.ANALYZE_RESPONSE,
+                success=False,
+                error="No previous response to analyze",
+            )
+
+        body = last_response.response_body.lower()
+
+        # Check for various indicators
+        found = False
+        extracted = {}
+
+        if check_for == "error":
+            error_keywords = ["error", "exception", "warning", "syntax", "unexpected"]
+            found = any(kw in body for kw in error_keywords)
+            extracted["errors_found"] = found
+
+        elif check_for == "success_indicator":
+            success_keywords = ["success", "welcome", "logged in", "dashboard"]
+            found = any(kw in body for kw in success_keywords)
+            extracted["success_found"] = found
+
+        elif check_for == "data":
+            # Look for data patterns
+            import re
+            emails = re.findall(r'[\w\.-]+@[\w\.-]+', last_response.response_body)
+            extracted["emails"] = emails[:10]
+            found = len(emails) > 0
+
+        return ActionResult(
+            action=ExploitAction.ANALYZE_RESPONSE,
+            success=found,
+            extracted_data=extracted,
+        )
+
+    async def _handle_extract_data(
+        self,
+        finding: Finding,
+        parameters: dict[str, Any],
+        last_response: Optional[ActionResult],
+    ) -> ActionResult:
+        """Handle extract_data action."""
+        import re
+
+        pattern = parameters.get("pattern", "")
+        data_type = parameters.get("data_type", "generic")
+
+        if not last_response or not last_response.response_body:
+            return ActionResult(
+                action=ExploitAction.EXTRACT_DATA,
+                success=False,
+                error="No response to extract from",
+            )
+
+        body = last_response.response_body
+        extracted = {}
+
+        try:
+            if pattern:
+                matches = re.findall(pattern, body)
+                extracted["matches"] = matches[:20]
+            else:
+                # Use data_type specific patterns
+                if data_type == "credentials":
+                    extracted["usernames"] = re.findall(r'username["\']?\s*[:=]\s*["\']?(\w+)', body, re.I)
+                    extracted["passwords"] = re.findall(r'password["\']?\s*[:=]\s*["\']?([^\s"\']+)', body, re.I)
+                elif data_type == "tokens":
+                    extracted["tokens"] = re.findall(r'token["\']?\s*[:=]\s*["\']?([a-zA-Z0-9_-]{20,})', body, re.I)
+                elif data_type == "database":
+                    extracted["tables"] = re.findall(r'table[:\s]+([a-zA-Z_]+)', body, re.I)
+                elif data_type == "files":
+                    extracted["paths"] = re.findall(r'(/[a-zA-Z0-9_/.-]+)', body)
+
+            success = any(v for v in extracted.values() if v)
+
+            return ActionResult(
+                action=ExploitAction.EXTRACT_DATA,
+                success=success,
+                extracted_data=extracted,
+            )
+
+        except re.error as e:
+            return ActionResult(
+                action=ExploitAction.EXTRACT_DATA,
+                success=False,
+                error=f"Regex error: {e}",
+            )
+
+    async def _handle_escalate(
+        self,
+        finding: Finding,
+        parameters: dict[str, Any],
+        last_response: Optional[ActionResult],
+    ) -> ActionResult:
+        """Handle escalation action."""
+        technique = parameters.get("technique", "")
+
+        # Escalation is context-dependent and may require custom implementation
+        return ActionResult(
+            action=ExploitAction.ESCALATE,
+            success=False,
+            error=f"Escalation technique '{technique}' requires custom implementation",
+        )
+
+    async def _handle_conclude(
+        self,
+        finding: Finding,
+        parameters: dict[str, Any],
+        last_response: Optional[ActionResult],
+    ) -> ActionResult:
+        """Handle conclusion action."""
+        success = parameters.get("success", False)
+        evidence = parameters.get("evidence", "")
+
+        return ActionResult(
+            action=ExploitAction.CONCLUDE,
+            success=success,
+            extracted_data={
+                "conclusion": "success" if success else "failure",
+                "evidence": evidence,
+            },
+        )
+
+    def _failed_result(
+        self,
+        finding: Finding,
+        reason: str,
+        steps: list[ExploitStep],
+        elapsed: float,
+        requests: int,
+    ) -> ExploitResult:
+        """Create a failed result."""
+        return ExploitResult(
+            success=False,
+            finding=finding,
+            steps=steps,
+            evidence=reason,
+            impact_achieved="None - exploitation failed",
+            exploitation_time_seconds=elapsed,
+            total_requests=requests,
+        )