"""
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,
        )
