#!/usr/bin/env python3
"""
Exploit suggestion engine.
Orchestrates version extraction, exploit matching, and ranking.
"""
from typing import Dict, List, Optional
from souleyez.intelligence.service_parser import ServiceVersionExtractor
from souleyez.intelligence.exploit_knowledge import ExploitKnowledgeBase
from souleyez.storage import exploit_attempts
import time


# Module-level cache shared across all instances
_SUGGESTION_CACHE = {}
_CACHE_TIMEOUT = 30


def _generate_exploit_identifier(exploit: Dict) -> str:
    """
    Generate a unique identifier for an exploit.

    Priority order:
    1. MSF module path (most specific)
    2. CVE ID
    3. EDB ID
    4. Title hash (fallback)

    Returns:
        Unique identifier string like "msf:exploit/unix/ftp/vsftpd_234_backdoor"
        or "cve:CVE-2011-2523" or "edb:49757" or "title:<hash>"
    """
    if exploit.get('msf_module'):
        return f"msf:{exploit['msf_module']}"
    elif exploit.get('cve'):
        return f"cve:{exploit['cve']}"
    elif exploit.get('edb_id'):
        return f"edb:{exploit['edb_id']}"
    else:
        # Fallback: use title (should be rare)
        title_hash = hash(exploit.get('title', 'unknown'))
        return f"title:{title_hash}"


class ExploitSuggestionEngine:
    """Generate exploit suggestions for engagement hosts."""

    def __init__(self, use_searchsploit: bool = True):
        """
        Initialize exploit suggestion engine.

        Args:
            use_searchsploit: Enable searchsploit integration (default: True)
        """
        self.version_extractor = ServiceVersionExtractor()
        self.knowledge_base = ExploitKnowledgeBase(use_searchsploit=use_searchsploit)
    
    def generate_suggestions(self, engagement_id: int, host_id: int = None) -> Dict:
        """
        Generate exploit suggestions for engagement or specific host.
        Results cached for 30 seconds to improve dashboard performance.
        
        Returns:
            {
                'hosts': [
                    {
                        'host_id': 123,
                        'ip': '10.0.0.5',
                        'hostname': 'example.com',
                        'services': [
                            {
                                'port': 21,
                                'service': 'ftp',
                                'version': 'vsftpd 2.3.4',
                                'product': 'vsftpd',
                                'version_number': '2.3.4',
                                'exploits': [
                                    {
                                        'title': 'vsftpd 2.3.4 Backdoor',
                                        'severity': 'critical',
                                        'msf_module': 'exploit/unix/ftp/vsftpd_234_backdoor',
                                        'cve': 'CVE-2011-2523',
                                        'description': '...',
                                        'match_type': 'exact'
                                    }
                                ],
                                'techniques': [
                                    'Anonymous FTP Access',
                                    'FTP Bounce Attack'
                                ]
                            }
                        ]
                    }
                ]
            }
        """
        # Check cache first
        cache_key = f"engagement_{engagement_id}_host_{host_id}"
        if cache_key in _SUGGESTION_CACHE:
            cached_result, cached_time = _SUGGESTION_CACHE[cache_key]
            if time.time() - cached_time < _CACHE_TIMEOUT:
                return cached_result
        
        # Cache miss or expired - do the analysis
        from souleyez.storage.hosts import HostManager
        from souleyez.storage.findings import FindingsManager
        
        hm = HostManager()
        fm = FindingsManager()
        
        # Get hosts to analyze
        if host_id:
            hosts = [hm.get_host(host_id)]
        else:
            hosts = hm.list_hosts(engagement_id)
        
        results = []
        
        for host in hosts:
            if not host:
                continue
            
            h_id = host.get('id')
            
            # Get findings for this host
            all_findings = fm.list_findings(engagement_id)
            host_findings = [
                f for f in all_findings 
                if f.get('host') == host.get('ip')
            ]
            
            # Extract versions from all sources
            services = self.version_extractor.combine_sources(h_id, host_findings)
            
            # Generate exploits for each service
            enriched_services = []
            for svc in services:
                exploits = self.knowledge_base.find_exploits(
                    svc['service'],
                    svc['version_number'],
                    svc.get('product')
                )

                # Enrich exploits with attempt status
                for exploit in exploits:
                    exploit_id = _generate_exploit_identifier(exploit)
                    exploit['identifier'] = exploit_id

                    # Look up attempt status
                    status = exploit_attempts.get_attempt_status(
                        engagement_id=engagement_id,
                        host_id=h_id,
                        exploit_identifier=exploit_id,
                        service_id=svc.get('service_id')
                    )
                    exploit['attempt_status'] = status or 'not_tried'

                techniques = self.knowledge_base.get_techniques(svc['service'])

                enriched_services.append({
                    **svc,
                    'exploits': exploits,
                    'techniques': techniques
                })
            
            # Only include hosts with exploits or techniques
            if any(s.get('exploits') or s.get('techniques') for s in enriched_services):
                results.append({
                    'host_id': h_id,
                    'ip': host.get('ip'),
                    'hostname': host.get('hostname'),
                    'services': enriched_services
                })
        
        result = {'hosts': results}
        
        # Cache the result
        _SUGGESTION_CACHE[cache_key] = (result, time.time())
        
        return result
    
    def get_host_summary(self, engagement_id: int, host_id: int) -> Dict:
        """
        Get exploit summary for a specific host.
        
        Returns:
            {
                'total_exploits': 5,
                'critical': 2,
                'high': 2,
                'medium': 1,
                'services_with_exploits': 3,
                'top_exploits': [...]
            }
        """
        suggestions = self.generate_suggestions(engagement_id, host_id)
        
        if not suggestions['hosts']:
            return {
                'total_exploits': 0,
                'critical': 0,
                'high': 0,
                'medium': 0,
                'low': 0,
                'services_with_exploits': 0,
                'top_exploits': []
            }
        
        host_data = suggestions['hosts'][0]
        all_exploits = []
        services_with_exploits = 0
        
        for svc in host_data['services']:
            if svc.get('exploits'):
                services_with_exploits += 1
                all_exploits.extend(svc['exploits'])
        
        # Count by severity
        severity_counts = {
            'critical': len([e for e in all_exploits if e.get('severity') == 'critical']),
            'high': len([e for e in all_exploits if e.get('severity') == 'high']),
            'medium': len([e for e in all_exploits if e.get('severity') == 'medium']),
            'low': len([e for e in all_exploits if e.get('severity') == 'low'])
        }
        
        # Get top 5 exploits
        top_exploits = sorted(
            all_exploits,
            key=lambda x: (
                {'critical': 4, 'high': 3, 'medium': 2, 'low': 1}.get(x.get('severity', 'low'), 0),
                {'exact': 3, 'range': 2, 'generic': 1}.get(x.get('match_type', 'generic'), 0)
            ),
            reverse=True
        )[:5]
        
        return {
            'total_exploits': len(all_exploits),
            **severity_counts,
            'services_with_exploits': services_with_exploits,
            'top_exploits': top_exploits
        }
    
    def get_engagement_summary(self, engagement_id: int) -> Dict:
        """
        Get exploit summary for entire engagement.
        
        Returns:
            {
                'total_hosts': 5,
                'hosts_with_exploits': 3,
                'total_exploits': 25,
                'critical': 5,
                'high': 10,
                'top_vulnerable_hosts': [...]
            }
        """
        suggestions = self.generate_suggestions(engagement_id)
        
        total_hosts = len(suggestions['hosts'])
        all_exploits = []
        
        host_exploit_counts = []
        
        for host in suggestions['hosts']:
            host_exploits = []
            for svc in host['services']:
                if svc.get('exploits'):
                    host_exploits.extend(svc['exploits'])
            
            all_exploits.extend(host_exploits)
            
            if host_exploits:
                host_exploit_counts.append({
                    'host_id': host['host_id'],
                    'ip': host['ip'],
                    'hostname': host.get('hostname'),
                    'exploit_count': len(host_exploits),
                    'critical_count': len([e for e in host_exploits if e.get('severity') == 'critical'])
                })
        
        # Sort hosts by exploit count
        top_vulnerable = sorted(
            host_exploit_counts,
            key=lambda x: (x['critical_count'], x['exploit_count']),
            reverse=True
        )[:10]
        
        return {
            'total_hosts': total_hosts,
            'hosts_with_exploits': len(host_exploit_counts),
            'total_exploits': len(all_exploits),
            'critical': len([e for e in all_exploits if e.get('severity') == 'critical']),
            'high': len([e for e in all_exploits if e.get('severity') == 'high']),
            'medium': len([e for e in all_exploits if e.get('severity') == 'medium']),
            'low': len([e for e in all_exploits if e.get('severity') == 'low']),
            'top_vulnerable_hosts': top_vulnerable
        }
