#!/usr/bin/env python3

"""
BRS-XSS Vulnerability Classifier

ML classifier for vulnerability type and severity assessment.

Company: EasyProTech LLC (www.easypro.tech)
Dev: Brabus
Modified: Sat 02 Aug 2025 09:35:54 MSK
Telegram: https://t.me/EasyProTech
"""

from typing import Dict, Tuple, Any

from .model_types import ModelType
from ..utils.logger import Logger

logger = Logger("ml.vulnerability_classifier")

try:
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False


class VulnerabilityClassifier:
    """
    ML vulnerability classifier.
    
    Determines vulnerability type and severity.
    """
    
    def __init__(self):
        """Initialize vulnerability classifier"""
        self.severity_model = None
        self.type_model = None
        self.is_trained = False
    
    def predict_severity(self, vulnerability_data: Dict[str, Any]) -> Tuple[str, float]:
        """
        Predict vulnerability severity.
        
        Args:
            vulnerability_data: Vulnerability data
            
        Returns:
            Tuple[severity_level, confidence]
        """
        # Simple heuristic evaluation
        context = vulnerability_data.get('context', '')
        payload = vulnerability_data.get('payload', '')
        reflection_accuracy = vulnerability_data.get('reflection_accuracy', 0.0)
        
        score = reflection_accuracy
        
        # Increase severity for dangerous contexts
        if 'javascript' in context.lower():
            score += 0.2
        if 'script' in payload.lower():
            score += 0.2
        if any(func in payload.lower() for func in ['alert', 'eval', 'function']):
            score += 0.1
        
        if score >= 0.8:
            return 'critical', 0.9
        elif score >= 0.6:
            return 'high', 0.8
        elif score >= 0.4:
            return 'medium', 0.7
        else:
            return 'low', 0.6
    
    def get_classifier_stats(self) -> Dict[str, Any]:
        """Classifier statistics"""
        return {
            'sklearn_available': SKLEARN_AVAILABLE,
            'models_trained': self.is_trained,
            'supported_model_types': [model_type.value for model_type in ModelType]
        }