"""
PROMETHEUS - EXPLOIT DETECTOR WITH LOCATION TRACKING

Enhanced to show exactly WHERE exploit patterns are found.

Copyright (c) 2026 Damian Donahue
"""

from typing import List, Dict, Any
from .models import ExploitMatch, Location, Severity
import re


class ExploitDetector:
    """
    Exploit pattern detection with location tracking.
    
    Detects common exploitation techniques and patterns.
    Now tracks EXACTLY where each pattern is found.
    """
    
    def __init__(self, intel_db: Dict[str, Any]):
        """Initialize with intelligence database."""
        self.patterns = intel_db.get('exploit_patterns', [])
    
    def detect(self, data: bytes) -> List[ExploitMatch]:
        """
        Detect exploit patterns and return matches with locations.
        
        Args:
            data: Binary data to analyze
            
        Returns:
            List of ExploitMatch objects with location data
        """
        matches = []
        
        # Detect NOP sleds (long runs of 0x90)
        matches.extend(self._detect_nop_sleds(data))
        
        # Detect other patterns from intelligence database
        for pattern in self.patterns:
            technique = pattern.get('technique', 'Unknown')
            pattern_type = pattern.get('pattern_type', 'unknown')
            signature = pattern.get('signature', '')
            
            # Skip NOP sleds (already detected above)
            if 'nop' in technique.lower():
                continue
            
            # Try to detect this pattern
            # This is simplified - real implementation would parse regex patterns
            if signature and len(signature) > 2:
                # For now, just look for simple byte sequences
                try:
                    # Convert hex pattern to bytes if applicable
                    if all(c in '0123456789abcdefABCDEFx{},' for c in signature):
                        # This looks like a hex pattern
                        # Simplified detection - real version would compile regex
                        pass
                except:
                    pass
        
        return matches
    
    def _detect_nop_sleds(self, data: bytes) -> List[ExploitMatch]:
        """
        Detect NOP sleds (sequences of 0x90 bytes).
        
        Returns matches with exact locations.
        """
        matches = []
        nop_runs = self._find_runs(data, 0x90, min_length=20)
        
        for offset, length in nop_runs:
            # Get context around NOP sled
            context_start = max(0, offset - 8)
            context_end = min(len(data), offset + length + 8)
            context_bytes = data[context_start:context_end]
            context_hex = ' '.join(f'{b:02x}' for b in context_bytes[:32])
            
            location = Location(
                offset=offset,
                length=length,
                context=f"... {context_hex} ..."
            )
            
            # Find the pattern in database for metadata
            pattern_meta = next(
                (p for p in self.patterns if 'nop' in p.get('technique', '').lower()),
                {}
            )
            
            match = ExploitMatch(
                technique=pattern_meta.get('technique', 'NOP Sled'),
                pattern_type=pattern_meta.get('pattern_type', 'buffer_overflow'),
                location=location,
                severity=Severity(pattern_meta.get('severity', 'high')),
                confidence=pattern_meta.get('confidence_weight', 0.8),
                explanation=pattern_meta.get('explanation', 'NOP sled detected'),
                context=pattern_meta.get('context', 'Shellcode preparation'),
                ttp_category=pattern_meta.get('ttp_category', 'Execution'),
                commonly_found_in=pattern_meta.get('commonly_found_in', [])
            )
            
            matches.append(match)
        
        return matches
    
    def _find_runs(self, data: bytes, byte: int, min_length: int) -> List[tuple]:
        """
        Find runs of a specific byte value.
        
        Args:
            data: Binary data to search
            byte: Byte value to look for
            min_length: Minimum run length
            
        Returns:
            List of (offset, length) tuples
        """
        runs = []
        current_offset = None
        current_length = 0
        
        for i, b in enumerate(data):
            if b == byte:
                if current_offset is None:
                    current_offset = i
                    current_length = 1
                else:
                    current_length += 1
            else:
                if current_length >= min_length:
                    runs.append((current_offset, current_length))
                current_offset = None
                current_length = 0
        
        # Handle final run
        if current_length >= min_length:
            runs.append((current_offset, current_length))
        
        return runs
