"""
Payload injection system for LogicPwn Exploit Engine.
- Supports static, dynamic, and template payloads
- Injects into URL, headers, body, JSON path, and form data
- Extensible and interoperable with core models
"""
from typing import List, Dict, Any
from logicpwn.models.request_config import RequestConfig
from .models import PayloadInjectionPoint
import copy
import json
import base64
import urllib.parse
import random
import string
import uuid

try:
    from jinja2 import Template
except ImportError:
    Template = None

class PayloadGenerator:
    @staticmethod
    def random_string(length=8):
        return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

    @staticmethod
    def random_int(min_val=0, max_val=10000):
        return random.randint(min_val, max_val)

    @staticmethod
    def random_uuid():
        return str(uuid.uuid4())

    @staticmethod
    def render_template(template_str: str, context: dict) -> str:
        if Template:
            return Template(template_str).render(**context)
        else:
            return template_str.format(**context)


def inject_payload(
    request_config: RequestConfig,
    injection_points: List[PayloadInjectionPoint],
    session_state: Dict[str, Any]
) -> RequestConfig:
    """
    Inject payloads into request at specified locations.
    Supports injection into:
    - URL parameters and path
    - HTTP headers
    - JSON body fields (using dot notation)
    - Form data
    - Raw body content
    """
    config = copy.deepcopy(request_config)
    for point in injection_points:
        value = point.payload_value
        # Advanced payload types
        if point.payload_type.value == "random":
            if point.target_field.lower().endswith("id"):
                value = PayloadGenerator.random_int()
            elif point.target_field.lower().endswith("uuid"):
                value = PayloadGenerator.random_uuid()
            else:
                value = PayloadGenerator.random_string()
        elif point.payload_type.value == "fuzz":
            # Simple fuzz: random string with special chars
            fuzz_chars = string.ascii_letters + string.digits + "'\"<>!@#$%^&*()"
            value = ''.join(random.choices(fuzz_chars, k=12))
        elif point.payload_type.value == "template":
            value = PayloadGenerator.render_template(value, session_state)
        # Encoding
        if point.encoding == "url":
            value = urllib.parse.quote_plus(value)
        elif point.encoding == "base64":
            value = base64.b64encode(value.encode()).decode()
        elif point.encoding == "html":
            import html
            value = html.escape(value)
        # Injection
        if point.location == "url":
            config.url = config.url.replace(f"{{{point.target_field}}}", value)
        elif point.location == "headers":
            config.headers = config.headers or {}
            config.headers[point.target_field] = value
        elif point.location == "body":
            config.raw_body = value
        elif point.location == "json_path":
            if not config.json_data:
                config.json_data = {}
            _set_json_path(config.json_data, point.target_field, value)
        elif point.location == "form_data":
            config.data = config.data or {}
            config.data[point.target_field] = value
    return config

def _set_json_path(obj: Dict[str, Any], path: str, value: Any):
    """
    Set a value in a nested dict using dot notation path (e.g., 'user.id').
    """
    keys = path.split('.')
    for k in keys[:-1]:
        if k not in obj or not isinstance(obj[k], dict):
            obj[k] = {}
        obj = obj[k]
    obj[keys[-1]] = value 