"""
Dynamic Testing Engine

Generates attack payloads and executes them in sandboxed test/staging
environments to confirm vulnerabilities identified during static analysis.

⚠️ WARNING: This module performs ACTIVE penetration testing.
Only use in authorized test environments.
"""

import asyncio
import time
import re
import logging
from typing import List, Dict, Optional, Any
from dataclasses import dataclass
from datetime import datetime
from enum import Enum


logger = logging.getLogger(__name__)


class TestStatus(Enum):
    """Dynamic test execution status"""
    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    TIMEOUT = "timeout"
    HALTED = "halted"


@dataclass
class TestResult:
    """Result of a single dynamic test"""
    test_id: str
    payload: str
    vulnerable: bool
    evidence: Dict[str, Any]
    status: TestStatus
    timestamp: str
    elapsed_time: float
    caused_error: bool = False


@dataclass
class AttackChainResult:
    """Result of multi-step attack chain"""
    chain_id: str
    steps_completed: int
    success: bool
    results: List[TestResult]
    total_time: float


class SecurityError(Exception):
    """Raised when security checks fail"""
    pass


class RateLimiter:
    """Rate limiting for dynamic testing"""
    
    def __init__(self, max_rps: float = 2.0):
        """
        Initialize rate limiter.
        
        Args:
            max_rps: Maximum requests per second
        """
        self.max_rps = max_rps
        self.min_interval = 1.0 / max_rps
        self.last_request_time = 0.0
    
    async def acquire(self):
        """Acquire permission to make a request (rate limited)"""
        now = time.time()
        time_since_last = now - self.last_request_time
        
        if time_since_last < self.min_interval:
            sleep_time = self.min_interval - time_since_last
            await asyncio.sleep(sleep_time)
        
        self.last_request_time = time.time()


class CircuitBreaker:
    """Circuit breaker to halt testing on errors"""
    
    def __init__(self, failure_threshold: int = 3):
        """
        Initialize circuit breaker.
        
        Args:
            failure_threshold: Number of failures before opening circuit
        """
        self.failure_threshold = failure_threshold
        self.failure_count = 0
        self.state = 'CLOSED'  # CLOSED, OPEN, HALF_OPEN
    
    def record_failure(self):
        """Record a failure"""
        self.failure_count += 1
        if self.failure_count >= self.failure_threshold:
            self.state = 'OPEN'
            logger.critical(f"Circuit breaker opened after {self.failure_count} failures")
    
    def record_success(self):
        """Record a success"""
        self.failure_count = max(0, self.failure_count - 1)
    
    def is_open(self) -> bool:
        """Check if circuit breaker is open"""
        return self.state == 'OPEN'
    
    def reset(self):
        """Reset circuit breaker"""
        self.failure_count = 0
        self.state = 'CLOSED'


class DynamicTestingEngine:
    """
    Generates attack payloads and executes them in sandboxed
    test/staging environments to confirm vulnerabilities.
    
    ⚠️ CRITICAL SAFETY: Always validates target environment before execution.
    """
    
    def __init__(self, test_environment_config: Dict):
        """
        Initialize dynamic testing engine.
        
        Args:
            test_environment_config: Configuration for test environment
                {
                    'base_url': 'https://staging.example.com',
                    'environment': 'staging',
                    'authorization_token': 'token',
                    'max_rps': 2,
                    'timeout': 30,
                    'halt_on_error': True
                }
        """
        self.env = test_environment_config
        self.rate_limiter = RateLimiter(max_rps=test_environment_config.get('max_rps', 2))
        self.circuit_breaker = CircuitBreaker(failure_threshold=3)
        self.audit_log = []
        
        # Validate environment is safe for testing
        self._validate_test_environment()
    
    def _validate_test_environment(self):
        """
        Validate that target environment is safe for active testing.
        Raises SecurityError if validation fails.
        """
        target_url = self.env.get('base_url', '')
        
        # Check 1: Not production
        forbidden_patterns = [
            r'prod(?:uction)?(?:\.|$)',
            r'api\.[a-z]+\.com$',
            r'www\.',
        ]
        
        for pattern in forbidden_patterns:
            if re.search(pattern, target_url, re.IGNORECASE):
                raise SecurityError(
                    f"Target appears to be production environment: {target_url}. "
                    f"Dynamic testing is prohibited in production."
                )
        
        # Check 2: Staging/test indicators
        safe_patterns = [
            r'staging',
            r'test',
            r'dev',
            r'localhost',
            r'127\.0\.0\.1',
            r'\.local$',
        ]
        
        if not any(re.search(p, target_url, re.IGNORECASE) for p in safe_patterns):
            raise SecurityError(
                f"Cannot verify target is test environment: {target_url}. "
                f"URL must contain: staging, test, dev, localhost, or .local"
            )
        
        # Check 3: Authorization token present
        if not self.env.get('authorization_token'):
            raise SecurityError(
                "Missing authorization token for dynamic testing. "
                "Set PENTEST_AUTHORIZATION_TOKEN environment variable."
            )
        
        logger.info(f"✅ Environment validation passed: {target_url}")
    
    async def generate_test_plan(self, sast_findings: List) -> List[Dict]:
        """
        Convert SAST findings into executable attack scenarios.
        
        Args:
            sast_findings: List of static analysis findings
        
        Returns:
            List of TestScenario objects
        """
        test_scenarios = []
        
        for finding in sast_findings:
            # Only generate dynamic tests for high-priority findings
            if finding.severity not in ['CRITICAL', 'HIGH']:
                continue
            
            if 'sql' in finding.category.lower() and 'injection' in finding.category.lower():
                scenario = {
                    "id": f"DYN-{finding.id}",
                    "type": "sql_injection",
                    "priority": "high",
                    "target": self._extract_endpoint_from_finding(finding),
                    "vulnerable_parameter": self._extract_parameter_from_finding(finding),
                    "payloads": self._get_sql_injection_payloads(),
                    "confirmation_criteria": {
                        "success_indicators": [
                            "status_code: 200",
                            "response_time_normal: true"
                        ],
                        "vulnerability_patterns": [
                            r"sql|mysql|postgres|syntax error",
                            r"you have an error in your sql syntax"
                        ]
                    }
                }
                test_scenarios.append(scenario)
            
            elif 'xss' in finding.category.lower():
                scenario = {
                    "id": f"DYN-{finding.id}",
                    "type": "xss",
                    "priority": "medium",
                    "target": self._extract_endpoint_from_finding(finding),
                    "vulnerable_parameter": self._extract_parameter_from_finding(finding),
                    "payloads": self._get_xss_payloads(),
                    "confirmation_criteria": {
                        "success_indicators": [
                            "response_contains: <script>"
                        ],
                        "vulnerability_patterns": [
                            r"<script>.*</script>",
                            r"onerror=",
                            r"javascript:"
                        ]
                    }
                }
                test_scenarios.append(scenario)
        
        return test_scenarios
    
    async def execute_test_scenario(self, scenario: Dict) -> List[TestResult]:
        """
        Execute a single test scenario with all payloads.
        
        Args:
            scenario: Test scenario definition
        
        Returns:
            List of TestResult objects
        """
        logger.info(f"▶️  Executing test scenario: {scenario['id']} ({scenario['type']})")
        
        # Log test initiation
        self._audit_log({
            "action": "test_started",
            "scenario_id": scenario['id'],
            "type": scenario['type'],
            "target": scenario['target'],
            "timestamp": datetime.utcnow().isoformat()
        })
        
        results = []
        
        for payload in scenario['payloads']:
            # Rate limiting
            await self.rate_limiter.acquire()
            
            # Circuit breaker check
            if self.circuit_breaker.is_open():
                logger.warning("⚠️  Circuit breaker open - halting tests")
                self._audit_log({
                    "action": "test_halted",
                    "reason": "circuit_breaker_open"
                })
                break
            
            try:
                result = await self._execute_single_test(scenario, payload)
                results.append(result)
                
                # Update circuit breaker
                if result.caused_error:
                    self.circuit_breaker.record_failure()
                else:
                    self.circuit_breaker.record_success()
                
            except Exception as e:
                logger.error(f"❌ Test execution error: {e}")
                self.circuit_breaker.record_failure()
                self._audit_log({
                    "action": "test_error",
                    "error": str(e),
                    "payload": payload
                })
        
        # Log completion
        vulnerable_count = len([r for r in results if r.vulnerable])
        logger.info(f"✅ Scenario complete: {len(results)} tests, {vulnerable_count} vulnerabilities confirmed")
        
        return results
    
    async def _execute_single_test(
        self,
        scenario: Dict,
        payload: str
    ) -> TestResult:
        """
        Execute a single test with a specific payload.
        
        Args:
            scenario: Test scenario
            payload: Attack payload
        
        Returns:
            TestResult object
        """
        test_id = f"{scenario['id']}-{hash(payload) % 10000:04d}"
        start_time = time.time()
        
        # Simulate HTTP request (in real implementation, use aiohttp)
        # For MVP, we'll simulate the test execution
        logger.debug(f"  Testing payload: {payload[:50]}...")
        
        try:
            # Simulate request execution
            await asyncio.sleep(0.1)  # Simulate network delay
            
            # Simulated response
            response = self._simulate_response(scenario, payload)
            
            elapsed_time = time.time() - start_time
            
            # Check vulnerability indicators
            is_vulnerable = self._check_vulnerability_indicators(
                response,
                scenario['confirmation_criteria']
            )
            
            # Build evidence
            evidence = {
                "request": {
                    "target": scenario['target'],
                    "parameter": scenario.get('vulnerable_parameter'),
                    "payload": payload,
                    "timestamp": datetime.utcnow().isoformat()
                },
                "response": {
                    "status_code": response['status_code'],
                    "body_snippet": response['body'][:500],
                    "elapsed_time": elapsed_time
                }
            }
            
            result = TestResult(
                test_id=test_id,
                payload=payload,
                vulnerable=is_vulnerable,
                evidence=evidence,
                status=TestStatus.SUCCESS,
                timestamp=datetime.utcnow().isoformat(),
                elapsed_time=elapsed_time,
                caused_error=(response['status_code'] >= 500)
            )
            
            if is_vulnerable:
                logger.warning(f"  🔴 VULNERABLE: {payload[:50]}...")
            
            return result
            
        except asyncio.TimeoutError:
            return TestResult(
                test_id=test_id,
                payload=payload,
                vulnerable=False,
                evidence={"timeout": True},
                status=TestStatus.TIMEOUT,
                timestamp=datetime.utcnow().isoformat(),
                elapsed_time=time.time() - start_time
            )
    
    def _simulate_response(self, scenario: Dict, payload: str) -> Dict:
        """
        Simulate HTTP response for testing.
        In production, this would use actual HTTP client (aiohttp).
        """
        
        # Simulate different responses based on payload
        if "' OR '1'='1" in payload or "admin' --" in payload:
            # Simulated SQL injection vulnerability
            return {
                "status_code": 200,
                "body": "Welcome admin! You have an error in your SQL syntax"
            }
        elif "<script>" in payload:
            # Simulated XSS vulnerability
            return {
                "status_code": 200,
                "body": f"Search results for: {payload}"
            }
        else:
            # Normal response
            return {
                "status_code": 200,
                "body": "Normal response"
            }
    
    def _check_vulnerability_indicators(
        self,
        response: Dict,
        confirmation_criteria: Dict
    ) -> bool:
        """
        Determine if response indicates vulnerability.
        
        Args:
            response: HTTP response data
            confirmation_criteria: Criteria for confirming vulnerability
        
        Returns:
            True if vulnerability confirmed
        """
        
        # Check for vulnerability-specific patterns
        vuln_patterns = confirmation_criteria.get('vulnerability_patterns', [])
        
        for pattern in vuln_patterns:
            if re.search(pattern, response['body'], re.IGNORECASE):
                return True  # Definite vulnerability
        
        return False
    
    async def execute_attack_chain(
        self,
        chain: Dict
    ) -> AttackChainResult:
        """
        Execute multi-step attack chain.
        
        Args:
            chain: Attack chain definition
                {
                    'id': 'CHAIN-001',
                    'steps': [
                        {'type': 'auth_bypass', 'target': '...'},
                        {'type': 'idor', 'target': '...', 'requires': 'auth_token'}
                    ]
                }
        
        Returns:
            AttackChainResult
        """
        logger.info(f"🔗 Executing attack chain: {chain['id']}")
        
        start_time = time.time()
        state = {}  # Shared state between steps
        results = []
        
        for step in chain['steps']:
            # Check prerequisites
            if 'requires' in step:
                if step['requires'] not in state:
                    logger.warning(f"  ⏭️  Skipping step - missing: {step['requires']}")
                    break
            
            # Execute step (simulated)
            logger.info(f"  Step: {step['type']}")
            await asyncio.sleep(0.1)
            
            # Simulate step result
            step_result = TestResult(
                test_id=f"{chain['id']}-{step['type']}",
                payload="",
                vulnerable=True,  # Simulated
                evidence={"step": step['type']},
                status=TestStatus.SUCCESS,
                timestamp=datetime.utcnow().isoformat(),
                elapsed_time=0.1
            )
            
            results.append(step_result)
            
            # Update state (simulated)
            if step['type'] == 'auth_bypass':
                state['auth_token'] = 'simulated_token'
        
        total_time = time.time() - start_time
        
        return AttackChainResult(
            chain_id=chain['id'],
            steps_completed=len(results),
            success=all(r.status == TestStatus.SUCCESS for r in results),
            results=results,
            total_time=total_time
        )
    
    def _extract_endpoint_from_finding(self, finding) -> str:
        """Extract API endpoint from SAST finding"""
        # Simplified - in real implementation, parse from code
        if hasattr(finding, 'file'):
            if 'api' in finding.file:
                return '/api/endpoint'
        return '/endpoint'
    
    def _extract_parameter_from_finding(self, finding) -> str:
        """Extract vulnerable parameter from SAST finding"""
        # Simplified - in real implementation, parse from code
        if 'username' in finding.code_snippet:
            return 'username'
        return 'input'
    
    def _get_sql_injection_payloads(self) -> List[str]:
        """Get SQL injection test payloads"""
        return [
            "' OR '1'='1",
            "' OR '1'='1' --",
            "admin' --",
        ]
    
    def _get_xss_payloads(self) -> List[str]:
        """Get XSS test payloads"""
        return [
            "<script>alert('XSS')</script>",
            "<img src=x onerror=alert('XSS')>",
        ]
    
    def _audit_log(self, entry: Dict):
        """Add entry to audit log"""
        self.audit_log.append({
            **entry,
            "timestamp": datetime.utcnow().isoformat()
        })
    
    def export_audit_log(self, output_path: str):
        """Export audit log to file"""
        import json
        with open(output_path, 'w') as f:
            json.dump(self.audit_log, f, indent=2)
        logger.info(f"📝 Audit log exported: {output_path}")

