vulnerability_scorer.py

from typing import Dict
import math
import logging

# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

class VulnerabilityScorer:
    def __init__(self):
        logger.debug("Initializing VulnerabilityScorer")
        self.cvss_weights = {
            'AV': {'N': 0.85, 'A': 0.62, 'L': 0.55, 'P': 0.2},
            'AC': {'L': 0.77, 'H': 0.44},
            'PR': {'N': 0.85, 'L': 0.62, 'H': 0.27},
            'UI': {'N': 0.85, 'R': 0.62},
            'S': {'U': 6.42, 'C': 7.52},
            'C': {'N': 0, 'L': 0.22, 'H': 0.56},
            'I': {'N': 0, 'L': 0.22, 'H': 0.56},
            'A': {'N': 0, 'L': 0.22, 'H': 0.56}
        }
        logger.debug(f"CVSS weights initialized: {self.cvss_weights}")

    def calculate_cvss_score(self, vector: str) -> float:
        logger.debug(f"Calculating CVSS score for vector: {vector}")
        metrics = dict(m.split(':') for m in vector.split('/')[1:])
        logger.debug(f"Parsed metrics: {metrics}")
        
        impact = 1 - ((1 - self.cvss_weights['C'][metrics['C']]) *
                      (1 - self.cvss_weights['I'][metrics['I']]) *
                      (1 - self.cvss_weights['A'][metrics['A']]))
        logger.debug(f"Calculated impact: {impact}")
        
        exploitability = (8.22 * self.cvss_weights['AV'][metrics['AV']] *
                          self.cvss_weights['AC'][metrics['AC']] *
                          self.cvss_weights['PR'][metrics['PR']] *
                          self.cvss_weights['UI'][metrics['UI']])
        logger.debug(f"Calculated exploitability: {exploitability}")
        
        if metrics['S'] == 'U':
            score = min((impact + exploitability), 10)
            logger.debug("Using unchanged scope calculation")
        else:
            score = min(1.08 * (impact + exploitability), 10)
            logger.debug("Using changed scope calculation")
        
        final_score = round(score, 1)
        logger.debug(f"Final CVSS score: {final_score}")
        return final_score

    def calculate_custom_score(self, vulnerability: Dict) -> float:
        logger.debug(f"Calculating custom score for vulnerability: {vulnerability}")
        base_score = self.calculate_cvss_score(vulnerability['cvss_vector'])
        logger.debug(f"Base CVSS score: {base_score}")
        
        time_factor = math.exp(-0.1 * vulnerability['days_since_disclosure'])
        logger.debug(f"Time factor: {time_factor}")
        
        exploit_factor = 1.5 if vulnerability['exploit_available'] else 1
        logger.debug(f"Exploit factor: {exploit_factor}")
        
        final_score = round(base_score * time_factor * exploit_factor, 2)
        logger.debug(f"Final custom score: {final_score}")
        return final_score

logger.debug("Creating VulnerabilityScorer instance")
scorer = VulnerabilityScorer()
logger.debug("VulnerabilityScorer instance created")