"""
XSS 漏洞检测
检测反射型和存储型 XSS 漏洞
"""
import re
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any, Set
from datetime import datetime
from playwright.async_api import Page
from loguru import logger


@dataclass
class XSSConfig:
    """XSS 检测配置"""
    timeout: float = 5000
    test_dom_xss: bool = True
    test_reflected_xss: bool = True
    check_event_handlers: bool = True
    check_innerhtml: bool = True
    payload_list: Optional[List[str]] = None


@dataclass
class XSSIssue:
    """XSS 漏洞"""
    severity: str  # high, medium, low
    type: str  # reflected, stored, dom-based
    location: str
    payload: str
    description: str
    evidence: Optional[str] = None
    recommendation: str = ""


@dataclass
class XSSScanResult:
    """XSS 扫描结果"""
    url: str
    issues: List[XSSIssue] = field(default_factory=list)
    scanned_inputs: int = 0
    test_payloads: int = 0
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())

    @property
    def has_issues(self) -> bool:
        return len(self.issues) > 0

    @property
    def high_severity_count(self) -> int:
        return sum(1 for i in self.issues if i.severity == "high")

    @property
    def medium_severity_count(self) -> int:
        return sum(1 for i in self.issues if i.severity == "medium")


class XSSDetector:
    """
    XSS 漏洞检测器
    检测常见 XSS 漏洞
    """

    # 常见 XSS Payload
    DEFAULT_PAYLOADS = [
        # 基础脚本注入
        "<script>alert('XSS')</script>",
        "<script>alert(1)</script>",
        "<img src=x onerror=alert(1)>",
        "<svg onload=alert(1)>",
        "<body onload=alert(1)>",
        # 事件处理器
        "\"onfocus=alert(1) autofocus=\"",
        "'onfocus=alert(1) autofocus='",
        "\"onmouseover=alert(1)\"",
        "'onclick=alert(1)'",
        # JavaScript 伪协议
        "javascript:alert(1)",
        "javascript:void(alert(1))",
        # 编码绕过
        "<script>alert(String.fromCharCode(88,83,83))</script>",
        "<img src=x onerror=\"alert&#40;1&#41;\">",
        # DOM XSS
        "#<script>alert(1)</script>",
        "?param=<script>alert(1)</script>",
    ]

    # 危险事件处理器
    DANGEROUS_EVENTS = [
        "onload", "onerror", "onclick", "onmouseover", "onfocus",
        "onblur", "onchange", "onsubmit", "onkeydown", "onkeyup",
        "ondrag", "ondrop", "onscroll", "onresize",
    ]

    # 危险 DOM 方法
    DANGEROUS_SINKS = [
        "innerHTML", "outerHTML", "document.write", "eval",
        "setTimeout", "setInterval", "jQuery.html",
    ]

    def __init__(
        self,
        page: Page,
        config: Optional[XSSConfig] = None,
    ):
        self.page = page
        self.config = config or XSSConfig()
        self._payloads = self.config.payload_list or self.DEFAULT_PAYLOADS

    # ==================== 扫描方法 ====================

    async def scan(self, url: Optional[str] = None) -> XSSScanResult:
        """
        执行 XSS 扫描

        Args:
            url: 要扫描的 URL（默认当前页面）

        Returns:
            XSSScanResult: 扫描结果
        """
        target_url = url or self.page.url
        result = XSSScanResult(url=target_url)

        # 检查事件处理器
        if self.config.check_event_handlers:
            event_issues = await self._check_event_handlers()
            result.issues.extend(event_issues)

        # 检查 innerHTML 使用
        if self.config.check_innerhtml:
            innerhtml_issues = await self._check_innerhtml_usage()
            result.issues.extend(innerhtml_issues)

        # 检查 DOM XSS 源和汇
        if self.config.test_dom_xss:
            dom_issues = await self._check_dom_xss()
            result.issues.extend(dom_issues)

        # 统计
        result.test_payloads = len(self._payloads)

        logger.info(f"XSS 扫描完成: {target_url}, 发现 {len(result.issues)} 个问题")
        return result

    async def test_input(
        self,
        selector: str,
        payload: Optional[str] = None,
    ) -> Optional[XSSIssue]:
        """
        测试输入点

        Args:
            selector: 输入元素选择器
            payload: 测试 Payload

        Returns:
            Optional[XSSIssue]: 发现的漏洞
        """
        test_payload = payload or self._payloads[0]

        try:
            # 获取元素
            element = self.page.locator(selector)
            if await element.count() == 0:
                return None

            # 注入 Payload
            await element.fill(test_payload)

            # 触发事件
            await element.press("Enter")

            # 检查是否执行
            xss_triggered = await self._check_xss_triggered()

            if xss_triggered:
                return XSSIssue(
                    severity="high",
                    type="reflected",
                    location=selector,
                    payload=test_payload,
                    description=f"输入点 '{selector}' 存在 XSS 漏洞",
                    recommendation="对用户输入进行 HTML 编码，使用安全的 DOM 操作方法",
                )

        except Exception as e:
            logger.debug(f"测试输入点失败: {selector}, error={e}")

        return None

    async def test_url_params(self, params: List[str]) -> List[XSSIssue]:
        """
        测试 URL 参数

        Args:
            params: 参数名列表

        Returns:
            List[XSSIssue]: 发现的漏洞
        """
        issues = []
        base_url = self.page.url.split("?")[0]

        for param in params:
            for payload in self._payloads[:5]:  # 限制测试数量
                test_url = f"{base_url}?{param}={payload}"

                try:
                    await self.page.goto(test_url)

                    # 检查 XSS 是否触发
                    if await self._check_xss_triggered():
                        issues.append(XSSIssue(
                            severity="high",
                            type="reflected",
                            location=f"URL 参数: {param}",
                            payload=payload,
                            description=f"URL 参数 '{param}' 存在反射型 XSS",
                            recommendation="对 URL 参数进行严格验证和编码",
                        ))
                        break

                except Exception as e:
                    logger.debug(f"测试 URL 参数失败: {param}, error={e}")

        return issues

    # ==================== 内部检查方法 ====================

    async def _check_event_handlers(self) -> List[XSSIssue]:
        """检查危险事件处理器"""
        issues = []

        for event in self.DANGEROUS_EVENTS:
            # 查找带有内联事件处理器的元素
            elements = self.page.locator(f"[{event}]")
            count = await elements.count()

            if count > 0:
                # 检查事件处理器内容
                for i in range(count):
                    element = elements.nth(i)
                    handler = await element.get_attribute(event)

                    if handler and self._is_dangerous_handler(handler):
                        issues.append(XSSIssue(
                            severity="medium",
                            type="dom-based",
                            location=f"元素 [{event}]",
                            payload=handler,
                            description=f"发现危险事件处理器: {event}",
                            evidence=handler[:100],
                            recommendation="避免使用内联事件处理器，使用 addEventListener",
                        ))

        return issues

    async def _check_innerhtml_usage(self) -> List[XSSIssue]:
        """检查 innerHTML 使用"""
        issues = []

        # 检查页面脚本中的 innerHTML 使用
        scripts = await self.page.evaluate("""
            () => {
                const scripts = Array.from(document.scripts);
                return scripts.map(s => s.textContent || s.innerHTML).join('\\n');
            }
        """)

        for sink in self.DANGEROUS_SINKS:
            if sink in scripts:
                issues.append(XSSIssue(
                    severity="medium",
                    type="dom-based",
                    location="页面脚本",
                    payload=sink,
                    description=f"检测到危险的 DOM 方法: {sink}",
                    recommendation="使用 textContent 替代 innerHTML，或对内容进行净化",
                ))

        return issues

    async def _check_dom_xss(self) -> List[XSSIssue]:
        """检查 DOM XSS"""
        issues = []

        # 检查 location.hash 是否直接用于 DOM 操作
        dom_issues = await self.page.evaluate("""
            () => {
                const issues = [];
                const scripts = document.querySelectorAll('script');

                scripts.forEach(script => {
                    const content = script.textContent || '';
                    // 检查是否直接使用 location 数据
                    if (content.includes('location.hash') ||
                        content.includes('location.search') ||
                        content.includes('location.href')) {
                        // 检查是否与危险 sink 结合
                        const sinks = ['innerHTML', 'document.write', 'eval'];
                        sinks.forEach(sink => {
                            if (content.includes(sink)) {
                                issues.push({
                                    type: 'dom-xss-source',
                                    sink: sink,
                                    description: 'location 数据可能被用于危险的 DOM 操作'
                                });
                            }
                        });
                    }
                });

                return issues;
            }
        """)

        for issue in dom_issues:
            issues.append(XSSIssue(
                severity="high",
                type="dom-based",
                location="页面脚本",
                payload=issue.get("sink", ""),
                description=issue.get("description", "潜在的 DOM XSS"),
                recommendation="不要直接将 location 数据用于 DOM 操作，必须先进行净化",
            ))

        return issues

    async def _check_xss_triggered(self) -> bool:
        """检查 XSS 是否被触发"""
        try:
            # 检查是否存在我们注入的标记
            alert_triggered = await self.page.evaluate("""
                () => {
                    // 检查是否有 alert 被触发（通过 window 属性标记）
                    return window.__xss_triggered === true;
                }
            """)
            return alert_triggered
        except Exception:
            return False

    def _is_dangerous_handler(self, handler: str) -> bool:
        """判断事件处理器是否危险"""
        dangerous_patterns = [
            r"eval\s*\(",
            r"document\.write",
            r"\.innerHTML\s*=",
            r"javascript:",
            r"alert\s*\(",
        ]

        for pattern in dangerous_patterns:
            if re.search(pattern, handler, re.IGNORECASE):
                return True
        return False

    # ==================== 便捷方法 ====================

    async def quick_scan(self) -> bool:
        """
        快速扫描，返回是否有漏洞

        Returns:
            bool: 是否发现漏洞
        """
        result = await self.scan()
        return result.has_issues

    async def get_event_handlers(self) -> Dict[str, List[str]]:
        """获取所有事件处理器"""
        handlers = {}

        for event in self.DANGEROUS_EVENTS:
            elements = self.page.locator(f"[{event}]")
            count = await elements.count()

            if count > 0:
                handlers[event] = []
                for i in range(count):
                    handler = await elements.nth(i).get_attribute(event)
                    if handler:
                        handlers[event].append(handler)

        return handlers


def create_xss_detector(
    page: Page,
    config: Optional[XSSConfig] = None,
) -> XSSDetector:
    """创建 XSS 检测器实例"""
    return XSSDetector(page, config)
