#!/usr/bin/env python3

"""
BRS-XSS Vulnerability Classifier

Classification of DOM XSS vulnerabilities by type and risk level.

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 Tuple

from .vulnerability_types import VulnerabilityType, RiskLevel
from .data_models import DataFlow
from .ast_types import ASTNode


class VulnerabilityClassifier:
    """Vulnerability classifier"""
    
    @staticmethod
    def classify_vulnerability(
        source_node: ASTNode, 
        sink_node: ASTNode, 
        data_flow: DataFlow
    ) -> Tuple[VulnerabilityType, RiskLevel]:
        """
        Classify vulnerability.
        
        Args:
            source_node: Source node
            sink_node: Sink node
            data_flow: Data flow
            
        Returns:
            Tuple[VulnerabilityType, RiskLevel]
        """
        
        vuln_type = VulnerabilityClassifier._determine_vulnerability_type(
            source_node, sink_node, data_flow
        )
        
        risk_level = VulnerabilityClassifier._determine_risk_level(
            vuln_type, source_node, sink_node, data_flow
        )
        
        return vuln_type, risk_level
    
    @staticmethod
    def _determine_vulnerability_type(
        source_node: ASTNode, 
        sink_node: ASTNode, 
        data_flow: DataFlow
    ) -> VulnerabilityType:
        """Determine vulnerability type"""
        
        source_value = source_node.value.lower()
        sink_value = sink_node.value.lower()
        
        # PostMessage XSS
        if 'postmessage' in source_value or 'message.data' in source_value:
            return VulnerabilityType.POSTMESSAGE_XSS
        
        # WebSocket XSS
        if 'websocket' in source_value:
            return VulnerabilityType.WEBSOCKET_XSS
        
        # Storage XSS
        if any(storage in source_value for storage in ['localstorage', 'sessionstorage']):
            return VulnerabilityType.STORAGE_XSS
        
        # URL manipulation
        if any(url_source in source_value for url_source in ['location', 'document.url']):
            if 'location' in sink_value:
                return VulnerabilityType.URL_MANIPULATION
        
        # Event handler
        if 'addeventlistener' in sink_value or any(event in sink_value for event in ['onclick', 'onload', 'onerror']):
            return VulnerabilityType.EVENT_HANDLER
        
        # Property injection
        if any(prop in sink_value for prop in ['innerhtml', 'outerhtml']):
            return VulnerabilityType.PROPERTY_INJECTION
        
        # Function parameter
        if sink_node.node_type.value == 'function_call':
            return VulnerabilityType.FUNCTION_PARAMETER
        
        # JSON injection
        if 'json' in source_value or 'json' in sink_value:
            return VulnerabilityType.JSON_INJECTION
        
        # Template injection
        if any(template in sink_value for template in ['template', 'render', 'compile']):
            return VulnerabilityType.TEMPLATE_INJECTION
        
        # Default - direct assignment
        return VulnerabilityType.DIRECT_ASSIGNMENT
    
    @staticmethod
    def _determine_risk_level(
        vuln_type: VulnerabilityType,
        source_node: ASTNode,
        sink_node: ASTNode,
        data_flow: DataFlow
    ) -> RiskLevel:
        """Determine risk level"""
        
        # Critical risk
        if vuln_type in [VulnerabilityType.DIRECT_ASSIGNMENT, VulnerabilityType.FUNCTION_PARAMETER]:
            if any(dangerous in sink_node.value.lower() for dangerous in ['eval', 'function', 'settimeout']):
                return RiskLevel.CRITICAL
        
        # High risk
        if vuln_type in [VulnerabilityType.PROPERTY_INJECTION, VulnerabilityType.EVENT_HANDLER]:
            return RiskLevel.HIGH
        
        if vuln_type == VulnerabilityType.URL_MANIPULATION:
            return RiskLevel.HIGH
        
        # Medium risk
        if vuln_type in [VulnerabilityType.POSTMESSAGE_XSS, VulnerabilityType.STORAGE_XSS]:
            return RiskLevel.MEDIUM
        
        # Low risk
        if vuln_type in [VulnerabilityType.WEBSOCKET_XSS, VulnerabilityType.JSON_INJECTION]:
            return RiskLevel.LOW
        
        # Consider sanitization
        if data_flow.has_sanitization and not data_flow.bypasses_sanitization:
            # Lower risk level
            risk_mapping = {
                RiskLevel.CRITICAL: RiskLevel.HIGH,
                RiskLevel.HIGH: RiskLevel.MEDIUM,
                RiskLevel.MEDIUM: RiskLevel.LOW,
                RiskLevel.LOW: RiskLevel.INFO
            }
            return risk_mapping.get(RiskLevel.HIGH, RiskLevel.INFO)  # Default fallback
        
        return RiskLevel.MEDIUM  # Default