#!/usr/bin/env python
"""
@Project ：django_base_ai
@File    ：xss_cleaned.py
@Author  ：congxing.wang
@Date    ：2026/02/25
@Desc    ：XSS 过滤与 HTML 清理
"""

import html as html_module
import re

import bleach
from bleach.css_sanitizer import CSSSanitizer

from django_base_ai.settings import ALLOWED_ATTRS, ALLOWED_STYLES, ALLOWED_TAGS

# 定义危险字符串和事件处理器列表
DANGEROUS_STRINGS = [
    "javascript:",
    "vbscript:",
    "data:",
    "onclick",
    "onerror",
    "onload",
    "onmouseover",
    "onfocus",
    "onblur",
    "onchange",
    "onsubmit",
    "onreset",
    "onselect",
    "onunload",
    "onkeydown",
    "onkeypress",
    "onkeyup",
    "onmousedown",
    "onmousemove",
    "onmouseout",
    "onmouseup",
    "onabort",
    "ondblclick",
    "onresize",
    "onscroll",
    "oncontextmenu",
    "oninput",
    "oninvalid",
    "onsearch",
    "onwheel",
    "oncopy",
    "oncut",
    "onpaste",
    "onbeforecopy",
    "onbeforecut",
    "onbeforepaste",
    "onhashchange",
    "onmessage",
    "onoffline",
    "ononline",
    "onpagehide",
    "onpageshow",
    "onpopstate",
    "onstorage",
    "onunhandledrejection",
    "onbeforeunload",
    "onclose",
    "onlanguagechange",
    "onorientationchange",
    "onvisibilitychange",
    "onbeforeprint",
    "onafterprint",
    "oncanplay",
    "oncanplaythrough",
    "ondurationchange",
    "onemptied",
    "onended",
    "onloadeddata",
    "onloadedmetadata",
    "onloadstart",
    "onpause",
    "onplay",
    "onplaying",
    "onprogress",
    "onratechange",
    "onseeked",
    "onseeking",
    "onstalled",
    "onsuspend",
    "ontimeupdate",
    "onvolumechange",
    "onwaiting",
    "onformchange",
    "onforminput",
    "oninvalid",
    "onsubmit",
    "onreset",
    "script",
    "iframe",
    "object",
    "embed",
    "form",
    "input",
    "textarea",
    "button",
    "select",
    "option",
    "meta",
    "link",
    "base",
    "applet",
    "frameset",
    "frame",
    "noscript",
    "isindex",
    "blink",
    "marquee",
]


def _remove_on_events(html):
    """
    移除所有HTML标签中的on事件属性（包括整个属性及其值）
    支持各种格式：onclick="...", onclick='...', onclick=..., ONCLICK="..."等
    处理所有边界情况，包括换行符、特殊字符、自闭合标签等
    :param html: 需要处理的HTML字符串
    :return: 处理后的HTML字符串
    """
    if not html:
        return html

    # 方法1：在HTML标签内部清理（最准确的方式）
    # 匹配所有HTML标签的开始部分（从<到>之间的内容）
    def clean_tag_attrs(match):
        """清理标签属性中的on事件"""
        tag_start = match.group(1)  # <tag
        tag_content = match.group(2)  # 属性部分
        tag_end = match.group(3)  # > 或 />

        if not tag_content:
            return match.group(0)

        # 移除所有on事件属性及其值
        # 使用更全面的正则表达式，处理各种情况
        cleaned = tag_content

        # 移除带双引号的on事件属性（处理前后可能有空格、换行符的情况）
        # 匹配：onxxx="..." 或 onclick="alert(1)" 等，包括多行内容
        cleaned = re.sub(r'\s*on\w+\s*=\s*"[^"]*"\s*', " ", cleaned, flags=re.IGNORECASE | re.DOTALL)
        # 移除带单引号的on事件属性（包括多行内容）
        cleaned = re.sub(r"\s*on\w+\s*=\s*'[^']*'\s*", " ", cleaned, flags=re.IGNORECASE | re.DOTALL)
        # 移除不带引号的on事件属性（值到下一个空格、>、换行符或引号为止）
        cleaned = re.sub(r'\s*on\w+\s*=\s*[^\s>"\'=]+', " ", cleaned, flags=re.IGNORECASE)

        # 清理多余的空格和换行符
        cleaned = re.sub(r"\s+", " ", cleaned).strip()

        # 重新组装
        if cleaned:
            return f"{tag_start} {cleaned}{tag_end}"
        else:
            return f"{tag_start}{tag_end}"

    # 匹配HTML标签：<tag ... > 或 <tag ... /> 或 <tag ...>
    # 使用非贪婪匹配确保只匹配到标签结束，支持多行属性
    pattern = r"(<[a-zA-Z][a-zA-Z0-9]*)\s*([^>]*?)(/?>)"
    html = re.sub(pattern, clean_tag_attrs, html, flags=re.IGNORECASE | re.DOTALL)

    # 方法2：直接在整个HTML中再次移除on事件属性（双重防护，处理可能遗漏的情况）
    # 处理带双引号的情况：onxxx="..."（包括多行）
    html = re.sub(r'\s+on\w+\s*=\s*"[^"]*"', "", html, flags=re.IGNORECASE | re.DOTALL)
    # 处理带单引号的情况：onxxx='...'（包括多行）
    html = re.sub(r"\s+on\w+\s*=\s*'[^']*'", "", html, flags=re.IGNORECASE | re.DOTALL)
    # 处理不带引号的情况：onxxx=value（值到下一个空格、>、引号或换行符为止）
    html = re.sub(r'\s+on\w+\s*=\s*[^\s>"\'=]+', "", html, flags=re.IGNORECASE)

    # 方法3：处理可能被编码或混淆的on事件（如onclick被拆分成on+click等）
    # 移除任何包含on事件模式的字符串（更激进的清理）
    html = re.sub(r'\s*on\w+\s*[:=]\s*["\'][^"\']*["\']', "", html, flags=re.IGNORECASE)
    html = re.sub(r'\s*on\w+\s*[:=]\s*[^\s>"\'=]+', "", html, flags=re.IGNORECASE)

    return html


def _filter_dangerous_strings(text):
    """
    过滤危险字符串和事件处理器
    :param text: 需要过滤的文本
    :return: 过滤后的文本
    """
    if not text:
        return text

    # 先移除所有on事件属性（更彻底的方式）
    text = _remove_on_events(text)

    # 过滤危险字符串（不区分大小写）
    for dangerous in DANGEROUS_STRINGS:
        # 使用正则表达式匹配，确保匹配完整的单词或协议
        pattern = re.compile(re.escape(dangerous), re.IGNORECASE)
        text = pattern.sub("", text)

    # 过滤 javascript: 和 vbscript: 协议（包括各种编码形式）
    text = re.sub(r"javascript\s*:", "", text, flags=re.IGNORECASE)
    text = re.sub(r"vbscript\s*:", "", text, flags=re.IGNORECASE)
    text = re.sub(r"data\s*:", "", text, flags=re.IGNORECASE)

    # 再次过滤事件处理器属性（on*），作为双重防护
    text = re.sub(r'\s+on\w+\s*=\s*["\'][^"\']*["\']', "", text, flags=re.IGNORECASE)
    text = re.sub(r"\s+on\w+\s*=\s*[^\s>]+", "", text, flags=re.IGNORECASE)

    # 过滤 <script> 标签及其内容
    text = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.IGNORECASE | re.DOTALL)
    text = re.sub(r"<script[^>]*>", "", text, flags=re.IGNORECASE)
    text = re.sub(r"</script>", "", text, flags=re.IGNORECASE)

    # 过滤其他危险标签
    dangerous_tags = ["iframe", "object", "embed", "form", "input", "textarea", "button", "select", "option"]
    for tag in dangerous_tags:
        text = re.sub(rf"<{tag}[^>]*>.*?</{tag}>", "", text, flags=re.IGNORECASE | re.DOTALL)
        text = re.sub(rf"<{tag}[^>]*>", "", text, flags=re.IGNORECASE)
        text = re.sub(rf"</{tag}>", "", text, flags=re.IGNORECASE)

    return text


def _preserve_user_mention_spans(html):
    """
    保护用户提及的span标签，这些标签包含特定的属性组合
    如果span标签中包含 style="color: #008AC5;"、userid 和 userinsert="true" 属性，则保留原格式
    :param html: 需要处理的HTML字符串
    :return: (处理后的HTML, 占位符到原始span的映射字典)
    """
    if not html:
        return html, {}

    # 匹配符合条件的span标签
    # 要求包含：style="color: #008AC5;"、userid="xxx"、userinsert="true"
    # 使用正则表达式匹配span标签及其内容，属性顺序可以任意
    # 先匹配所有span标签，然后检查是否包含必需的属性
    def is_valid_user_mention_span(span_tag):
        """检查span标签是否包含必需的属性"""
        # 检查是否包含 style="color: #008AC5;" (允许有或没有分号，不区分大小写)
        # 匹配 style="color: #008AC5;" 或 style="color:#008AC5;" 等变体
        has_style = re.search(r'style\s*=\s*["\'][^"\']*color\s*:\s*#008ac5\s*;?[^"\']*["\']', span_tag, re.IGNORECASE)
        # 检查是否包含 userid 属性（值可以是任意字符串）
        has_userid = re.search(r'userid\s*=\s*["\'][^"\']+["\']', span_tag, re.IGNORECASE)
        # 检查是否包含 userinsert="true"
        has_userinsert = re.search(r'userinsert\s*=\s*["\']true["\']', span_tag, re.IGNORECASE)
        return bool(has_style and has_userid and has_userinsert)

    # 匹配所有span标签（包括内容和闭合标签）
    pattern = r"<span\s+[^>]*>.*?</span>"

    placeholder_map = {}
    counter = 0

    def replace_span(match):
        nonlocal counter
        full_span = match.group(0)
        # 只保护符合条件的span标签
        if is_valid_user_mention_span(full_span):
            # 移除span标签中的on事件属性
            cleaned_span = _remove_on_events(full_span)
            placeholder = f"__PRESERVED_SPAN_{counter}__"
            placeholder_map[placeholder] = cleaned_span
            counter += 1
            return placeholder
        else:
            # 不符合条件的span标签保持原样，后续会被正常处理
            return full_span

    # 替换所有符合条件的span标签为占位符
    processed_html = re.sub(pattern, replace_span, html, flags=re.IGNORECASE | re.DOTALL)

    return processed_html, placeholder_map


def _restore_user_mention_spans(html, placeholder_map):
    """
    恢复被保护的span标签
    :param html: 处理后的HTML字符串
    :param placeholder_map: 占位符到原始span的映射字典
    :return: 恢复后的HTML字符串
    """
    if not placeholder_map:
        return html

    for placeholder, original_span in placeholder_map.items():
        html = html.replace(placeholder, original_span)

    return html


def _preserve_nbsp(html):
    """
    保护 &nbsp; 实体，避免被HTML转义
    匹配 &nbsp; 或 &NBSP; 等大小写变体，但不匹配已经被转义的 &amp;nbsp;
    :param html: 需要处理的HTML字符串
    :return: (处理后的HTML, 占位符到原始&nbsp;的映射字典)
    """
    if not html:
        return html, {}

    placeholder_map = {}
    counter = 0

    # 匹配 &nbsp; (不区分大小写)
    # 使用负向前瞻确保不是 &amp;nbsp; (已经被转义的)
    pattern = r"&(?!amp;)nbsp;"

    def replace_nbsp(match):
        nonlocal counter
        placeholder = f"__PRESERVED_NBSP_{counter}__"
        placeholder_map[placeholder] = "&nbsp;"
        counter += 1
        return placeholder

    # 替换所有 &nbsp; 为占位符（不区分大小写）
    processed_html = re.sub(pattern, replace_nbsp, html, flags=re.IGNORECASE)

    return processed_html, placeholder_map


def _restore_nbsp(html, placeholder_map):
    """
    恢复被保护的 &nbsp; 实体
    :param html: 处理后的HTML字符串
    :param placeholder_map: 占位符到原始&nbsp;的映射字典
    :return: 恢复后的HTML字符串
    """
    if not placeholder_map:
        return html

    for placeholder, original_nbsp in placeholder_map.items():
        html = html.replace(placeholder, original_nbsp)

    return html


def cleaned_string(html="", custome_allowed_styles=[], escape_html=True):
    """
    清理字符串，进行严格的XSS防护。
    对于评论等纯文本场景，应该完全转义HTML。
    特殊处理：保留包含用户提及信息的span标签（style="color: #008AC5;"、userid、userinsert="true"）
    :param html: 需要清理的字符串
    :param custome_allowed_styles: 自定义允许的样式列表
    :param escape_html: 是否对输出进行HTML转义（默认True，适用于纯文本场景）
    :return: 清理后的字符串
    """
    if not html:
        return html

    # 先保护符合条件的用户提及span标签
    html, span_placeholder_map = _preserve_user_mention_spans(html)

    # 保护 &nbsp; 实体，避免被转义
    html, nbsp_placeholder_map = _preserve_nbsp(html)

    # 先过滤危险字符串和事件处理器
    html = _filter_dangerous_strings(html)

    # 如果需要对HTML进行转义（纯文本场景），直接转义并返回
    if escape_html:
        # 对输出进行HTML转义，确保特殊字符被转义
        # 这是最安全的方式，适用于评论等纯文本场景
        cleaned_html = html_module.escape(html)
        # 再次过滤危险字符串（双重防护）
        cleaned_html = _filter_dangerous_strings(cleaned_html)
        # 恢复被保护的 &nbsp; 实体
        cleaned_html = _restore_nbsp(cleaned_html, nbsp_placeholder_map)
        # 恢复被保护的用户提及span标签
        cleaned_html = _restore_user_mention_spans(cleaned_html, span_placeholder_map)
        # 最后再次移除所有on事件，确保完全清除（包括恢复的span标签中的on事件）
        cleaned_html = _remove_on_events(cleaned_html)
        # 再次检查并移除on事件（三重防护，确保完全清除）
        cleaned_html = _remove_on_events(cleaned_html)
        return cleaned_html

    # 如果需要保留某些HTML标签（非纯文本场景），使用bleach进行清理
    # 实例化自定义的 CSS sanitizer
    if len(custome_allowed_styles) > 0:
        css_santizer = CSSSanitizer(allowed_css_properties=custome_allowed_styles)
    else:
        css_santizer = CSSSanitizer(allowed_css_properties=ALLOWED_STYLES)

    # 创建更严格的允许标签列表，确保移除所有危险标签
    safe_tags = [
        tag
        for tag in ALLOWED_TAGS
        if tag
        not in [
            "script",
            "object",
            "embed",
            "form",
            "input",
            "textarea",
            "button",
            "select",
            "option",
            "iframe",
            "meta",
            "link",
            "style",
            "base",
            "applet",
            "frameset",
            "frame",
            "noscript",
        ]
    ]

    # 创建更严格的属性配置，移除所有事件处理器
    safe_attrs = {}
    for tag, attrs in ALLOWED_ATTRS.items():
        # 过滤掉所有以 on 开头的属性（事件处理器）
        filtered_attrs = [attr for attr in attrs if not attr.lower().startswith("on")]
        if filtered_attrs:
            safe_attrs[tag] = filtered_attrs

    # 清理字符串，去除不需要的属性
    cleaned_html = bleach.clean(
        html, tags=safe_tags, attributes=safe_attrs, css_sanitizer=css_santizer, protocols=["http", "https", "mailto"]
    )

    # 再次过滤危险字符串（双重防护）
    cleaned_html = _filter_dangerous_strings(cleaned_html)

    # 恢复被保护的 &nbsp; 实体
    cleaned_html = _restore_nbsp(cleaned_html, nbsp_placeholder_map)
    # 恢复被保护的用户提及span标签
    cleaned_html = _restore_user_mention_spans(cleaned_html, span_placeholder_map)
    # 最后再次移除所有on事件，确保完全清除（包括恢复的span标签中的on事件）
    cleaned_html = _remove_on_events(cleaned_html)
    # 再次检查并移除on事件（三重防护，确保完全清除）
    cleaned_html = _remove_on_events(cleaned_html)

    return cleaned_html


def cleaned_html_file(html="", custome_allowed_styles=[]):
    """
    专门用于清理 HTML 文件的函数，移除所有危险标签和属性
    特殊处理：保留 &nbsp; 实体，避免被转义
    :param html: 需要清理的 HTML 字符串
    :param custome_allowed_styles: 自定义允许的样式列表
    :return: 清理后的字符串
    """
    if not html:
        return html

    # 保护 &nbsp; 实体，避免被转义
    html, nbsp_placeholder_map = _preserve_nbsp(html)

    # 先过滤危险字符串和事件处理器
    html = _filter_dangerous_strings(html)

    # 为 HTML 文件创建更严格的允许标签列表，移除危险标签
    safe_tags = [
        tag
        for tag in ALLOWED_TAGS
        if tag
        not in [
            "script",
            "object",
            "embed",
            "form",
            "input",
            "textarea",
            "button",
            "select",
            "option",
            "iframe",
            "meta",
            "link",
            "style",
            "base",
            "applet",
            "frameset",
            "frame",
            "noscript",
        ]
    ]

    # 为 HTML 文件创建更严格的属性配置，移除事件处理器
    safe_attrs = {
        "*": ["style", "class", "id"],  # 只允许基本的样式和标识属性
        "a": ["href", "title", "target"],  # a 标签允许的属性（过滤掉onclick等）
        "img": ["src", "alt", "width", "height", "title"],  # img 标签允许的属性
        "table": ["border", "cellpadding", "cellspacing", "width"],
        "td": ["colspan", "rowspan", "width", "height"],
        "th": ["colspan", "rowspan", "width", "height"],
    }

    # 过滤掉所有事件处理器属性
    for tag in safe_attrs:
        if isinstance(safe_attrs[tag], list):
            safe_attrs[tag] = [attr for attr in safe_attrs[tag] if not attr.lower().startswith("on")]

    # 实例化自定义的 CSS sanitizer
    if len(custome_allowed_styles) > 0:
        css_santizer = CSSSanitizer(allowed_css_properties=custome_allowed_styles)
    else:
        css_santizer = CSSSanitizer(allowed_css_properties=ALLOWED_STYLES)

    # 清理字符串，使用更严格的配置
    cleaned_html = bleach.clean(
        html, tags=safe_tags, attributes=safe_attrs, css_sanitizer=css_santizer, protocols=["http", "https", "mailto"]
    )

    # 再次过滤危险字符串（双重防护）
    cleaned_html = _filter_dangerous_strings(cleaned_html)

    # 恢复被保护的 &nbsp; 实体
    cleaned_html = _restore_nbsp(cleaned_html, nbsp_placeholder_map)

    # 最后再次移除所有on事件，确保完全清除（包括恢复的内容中的on事件）
    cleaned_html = _remove_on_events(cleaned_html)
    # 再次检查并移除on事件（双重防护，确保完全清除）
    cleaned_html = _remove_on_events(cleaned_html)

    # 返回处理后的 HTML 字符串
    return cleaned_html


if __name__ == "__main__":
    # 测试XSS防护功能
    test_cases = [
        "<script>alert(/XSS/);</script>",
        "<img src=1 onerror=alert(1) />",
        '<a href="javascript:alert(1)">Click</a>',
        '<p onclick="alert(1)">Test</p>',
        '<iframe src="javascript:alert(1)"></iframe>',
        "<p>正常文本</p>",
        "<script>alert(/XSS/;</script>",  # 用户报告的测试用例
        '<div onclick="alert(1)" onmouseover="alert(2)">测试</div>',  # 多个on事件
        "<span onclick='alert(1)'>测试</span>",  # 单引号
        '<p ONCLICK="alert(1)">大小写测试</p>',  # 大小写不敏感
        '<a href="#" onclick="alert(1)" onfocus="alert(2)">链接</a>',  # 多个on事件
        '<img src="test.jpg" onload="alert(1)" onerror="alert(2)" />',  # 自闭合标签
    ]

    print("=" * 60)
    print("XSS防护测试")
    print("=" * 60)

    for i, test_html in enumerate(test_cases, 1):
        print(f"\n测试用例 {i}: {test_html}")
        cleaned = cleaned_string(test_html, escape_html=True)
        print(f"清理后: {cleaned}")
        print(
            f"是否包含危险内容: {'是' if any(d in cleaned.lower() for d in ['script', 'onclick', 'javascript:', 'alert'] if d) else '否'}"
        )

    print("\n" + "=" * 60)
    print("用户提及span标签保留测试")
    print("=" * 60)

    # 测试用户提及span标签的保留功能
    user_mention_test = '34324<span style="color: #008AC5;" userid="134" userinsert="true">@测试用户5(TEST05)</span>&nbsp;234324 <h1>A<h1>内容<span style="color: #008AC5;" userid="134" userinsert="true">@测试用户5(TEST05)</span>&nbsp;内容'
    print(f"\n原始内容: {user_mention_test}")
    cleaned_mention = cleaned_string(user_mention_test, escape_html=True)
    print(f"清理后: {cleaned_mention}")
    # 检查是否保留了用户提及的span标签
    has_preserved_span = (
        '<span style="color: #008AC5;" userid="134" userinsert="true" onclick="alert(1)">' in cleaned_mention
    )
    print(f"是否保留了用户提及span标签: {'是' if has_preserved_span else '否'}")

    # 测试不符合条件的span标签应该被转义
    normal_span_test = '<span class="test">普通span</span>'
    print(f"\n普通span测试: {normal_span_test}")
    cleaned_normal = cleaned_string(normal_span_test, escape_html=True)
    print(f"清理后: {cleaned_normal}")
    # 普通span应该被转义
    is_escaped = "&lt;span" in cleaned_normal or "<span" not in cleaned_normal
    print(f"普通span是否被转义: {'是' if is_escaped else '否'}")

    print("\n" + "=" * 60)
    print("&nbsp; 保留测试")
    print("=" * 60)

    # 测试 &nbsp; 是否被正确保留
    nbsp_test = "测试&nbsp;内容&nbsp;更多内容"
    print(f"\n原始内容: {nbsp_test}")
    cleaned_nbsp = cleaned_string(nbsp_test, escape_html=True)
    print(f"清理后: {cleaned_nbsp}")
    # 检查 &nbsp; 是否被保留（不应该被转义为 &amp;nbsp;）
    has_nbsp = "&nbsp;" in cleaned_nbsp
    has_escaped_nbsp = "&amp;nbsp;" in cleaned_nbsp
    print(f"是否保留了 &nbsp;: {'是' if has_nbsp else '否'}")
    print(f"是否被转义为 &amp;nbsp;: {'是' if has_escaped_nbsp else '否'}")

    # 测试包含 &nbsp; 和用户提及span的混合内容
    mixed_test = '文本&nbsp;<span style="color: #008AC5;" userid="134" userinsert="true" onclick="alert(1)">@用户</span>&nbsp;更多文本'
    print(f"\n混合内容测试: {mixed_test}")
    cleaned_mixed = cleaned_string(mixed_test, escape_html=True)
    print(f"清理后: {cleaned_mixed}")
    has_nbsp_in_mixed = "&nbsp;" in cleaned_mixed
    has_span_in_mixed = '<span style="color: #008AC5;" userid="134" userinsert="true">' in cleaned_mixed
    print(f"是否保留了 &nbsp;: {'是' if has_nbsp_in_mixed else '否'}")
    print(f"是否保留了用户提及span: {'是' if has_span_in_mixed else '否'}")

    print("\n" + "=" * 60)
    print("on事件清除专项测试")
    print("=" * 60)

    # 专门测试on事件清除
    on_event_test_cases = [
        ('<p onclick="alert(1)">测试</p>', "onclick"),
        ('<div onmouseover="alert(2)">测试</div>', "onmouseover"),
        ('<a onclick="alert(1)" onfocus="alert(2)">链接</a>', "onclick"),
        ('<img onload="alert(1)" onerror="alert(2)" />', "onload"),
        ('<span ONCLICK="alert(1)">大小写</span>', "onclick"),
        ("<div onclick='alert(1)'>单引号</div>", "onclick"),
        ("<p onclick=alert(1)>无引号</p>", "onclick"),
        ('<div onclick="alert(1)" class="test">混合属性</div>', "onclick"),
    ]

    for test_html, event_name in on_event_test_cases:
        print(f"\n测试: {test_html}")
        cleaned = cleaned_string(test_html, escape_html=False)
        print(f"清理后: {cleaned}")
        has_event = event_name.lower() in cleaned.lower()
        print(f"是否包含{event_name}事件: {'是' if has_event else '否'}")
        if has_event:
            print("⚠️  警告：仍然包含on事件！")

    print("\n" + "=" * 60)
    print("测试完成")
    print("=" * 60)
