"""
YAML Metadata Generator for SIGMA Rules
Generates YAML metadata sections for SIGMA rules based on CVE and PoC data
"""

import logging
from typing import Dict, List, Optional, Any
from sqlalchemy.orm import Session
from datetime import datetime

logger = logging.getLogger(__name__)

class YAMLMetadataGenerator:
    """Generates YAML metadata sections for SIGMA rules"""
    
    def __init__(self, db_session: Session):
        self.db_session = db_session
    
    def generate_metadata(self, cve, poc_data: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Generate YAML metadata for a SIGMA rule based on CVE and PoC data
        
        Args:
            cve: CVE database object
            poc_data: List of PoC data dictionaries
            
        Returns:
            Dictionary containing YAML metadata sections
        """
        try:
            # Extract CVE information
            cve_id = cve.cve_id
            description = cve.description or ""
            cvss_score = getattr(cve, 'cvss_score', None)
            published_date = getattr(cve, 'published_date', None)
            
            # Determine attack techniques from PoC data
            attack_techniques = self._extract_attack_techniques(poc_data)
            
            # Generate basic metadata
            metadata = {
                'title': f"Potential {cve_id} Exploitation",
                'id': f"sigma-{cve_id.lower().replace('-', '_')}",
                'description': self._generate_description(cve_id, description),
                'references': [
                    f"https://cve.mitre.org/cgi-bin/cvename.cgi?name={cve_id}",
                    f"https://nvd.nist.gov/vuln/detail/{cve_id}"
                ],
                'author': "Auto-generated SIGMA Rule",
                'date': datetime.now().strftime("%Y/%m/%d"),
                'tags': self._generate_tags(cve_id, attack_techniques),
                'level': self._determine_level(cvss_score),
                'status': "experimental"
            }
            
            # Add PoC-specific references
            if poc_data:
                for poc in poc_data[:3]:  # Add up to 3 PoC references
                    if 'html_url' in poc:
                        metadata['references'].append(poc['html_url'])
            
            # Add MITRE ATT&CK techniques if available
            if attack_techniques:
                metadata['falsepositives'] = [
                    "Legitimate use of affected software",
                    "Administrative activities"
                ]
                metadata['fields'] = ["CommandLine", "ProcessName", "ParentProcessName"]
            
            return metadata
            
        except Exception as e:
            logger.error(f"Error generating metadata for {cve.cve_id}: {e}")
            return self._generate_fallback_metadata(cve.cve_id)
    
    def _extract_attack_techniques(self, poc_data: List[Dict[str, Any]]) -> List[str]:
        """Extract MITRE ATT&CK techniques from PoC data"""
        techniques = []
        
        for poc in poc_data:
            # Look for common attack patterns in PoC descriptions
            description = poc.get('description', '').lower()
            
            if 'remote code execution' in description or 'rce' in description:
                techniques.append('T1203')  # Exploitation for Client Execution
            if 'privilege escalation' in description:
                techniques.append('T1068')  # Exploitation for Privilege Escalation
            if 'sql injection' in description:
                techniques.append('T1190')  # Exploit Public-Facing Application
            if 'xss' in description or 'cross-site scripting' in description:
                techniques.append('T1185')  # Browser Session Hijacking
            if 'buffer overflow' in description:
                techniques.append('T1203')  # Exploitation for Client Execution
            if 'deserialization' in description:
                techniques.append('T1190')  # Exploit Public-Facing Application
        
        return list(set(techniques))
    
    def _generate_description(self, cve_id: str, description: str) -> str:
        """Generate a concise description for the SIGMA rule"""
        if description:
            # Take first sentence or first 200 characters
            first_sentence = description.split('.')[0]
            if len(first_sentence) > 200:
                return first_sentence[:200] + "..."
            return first_sentence + "."
        else:
            return f"Detects potential exploitation of {cve_id}"
    
    def _generate_tags(self, cve_id: str, attack_techniques: List[str]) -> List[str]:
        """Generate tags for the SIGMA rule"""
        tags = [
            "attack.t1203",  # Default to exploitation technique
            "cve." + cve_id.lower().replace('-', '_')
        ]
        
        # Add specific technique tags
        for technique in attack_techniques:
            tags.append(f"attack.{technique.lower()}")
        
        return tags
    
    def _determine_level(self, cvss_score: Optional[float]) -> str:
        """Determine the severity level based on CVSS score"""
        if cvss_score is None:
            return "medium"
        
        if cvss_score >= 9.0:
            return "critical"
        elif cvss_score >= 7.0:
            return "high"
        elif cvss_score >= 4.0:
            return "medium"
        else:
            return "low"
    
    def _generate_fallback_metadata(self, cve_id: str) -> Dict[str, Any]:
        """Generate minimal fallback metadata when primary generation fails"""
        return {
            'title': f"Potential {cve_id} Exploitation",
            'id': f"sigma-{cve_id.lower().replace('-', '_')}",
            'description': f"Detects potential exploitation of {cve_id}",
            'references': [
                f"https://cve.mitre.org/cgi-bin/cvename.cgi?name={cve_id}",
                f"https://nvd.nist.gov/vuln/detail/{cve_id}"
            ],
            'author': "Auto-generated SIGMA Rule",
            'date': datetime.now().strftime("%Y/%m/%d"),
            'tags': [
                "attack.t1203",
                f"cve.{cve_id.lower().replace('-', '_')}"
            ],
            'level': "medium",
            'status': "experimental"
        }