Risk Engine Module

class RiskEngine:
    def __init__(self):
        """Initialize risk engine"""
        self.risk_scores = {
            'CRITICAL': 10,
            'HIGH': 7,
            'MEDIUM': 4,
            'LOW': 2,
            'UNKNOWN': 1
        }
        
        self.severity_weights = {
            'exploitable': 2.0,
            'public_exploit': 1.5,
            'network_accessible': 1.3,
            'authenticated': 0.7
        }

def calculate_cvss_risk(self, cvss_score):
    """Convert CVSS score to risk level"""
    if cvss_score >= 9.0:
        return "CRITICAL"
    elif cvss_score >= 7.0:
        return "HIGH"
    elif cvss_score >= 4.0:
        return "MEDIUM"
    elif cvss_score > 0.0:
        return "LOW"
    else:
        return "UNKNOWN"

def calculate_port_risk(self, port_data, vulnerabilities, cves):
    """Calculate risk for a specific port"""
    risk_score = 0
    risk_factors = []
    
    port_num = port_data.get('port')
    service = port_data.get('service', '').lower()
    
    # Base risk for port exposure
    if port_num < 1024:  # Well-known ports
        risk_score += 1
        risk_factors.append("Well-known port exposed")
    
    # Check for sensitive services
    sensitive_services = ['ftp', 'telnet', 'smtp', 'mysql', 'postgresql', 
                         'microsoft-ds', 'netbios-ssn', 'rdp']
    
    if any(svc in service for svc in sensitive_services):
        risk_score += 2
        risk_factors.append(f"Sensitive service exposed: {service}")
    
    # Check vulnerabilities for this port
    port_vulns = [v for v in vulnerabilities if v.get('port') == port_num]
    
    for vuln in port_vulns:
        severity = vuln.get('severity', 'UNKNOWN')
        risk_score += self.risk_scores.get(severity, 1)
        risk_factors.append(f"{severity} vulnerability: {vuln.get('name')}")
        
        if vuln.get('exploitable', False):
            risk_score *= 1.5  # 50% increase for exploitable vulns
            risk_factors.append("Exploitable vulnerability detected")
    
    # Check CVEs for this port
    port_cves = [c for c in cves if c.get('port') == port_num]
    
    for cve in port_cves:
        cvss = cve.get('cvss_score', 0)
        risk_score += cvss / 2  # Weight CVE score
        
        if cve.get('exploitable', False):
            risk_score *= 1.3  # 30% increase for CVE with exploit
            risk_factors.append(f"CVE with public exploit: {cve.get('cve_id')}")
    
    # Normalize to 0-10 scale
    normalized_score = min(risk_score, 10.0)
    
    return {
        'port': port_num,
        'service': service,
        'risk_score': round(normalized_score, 2),
        'risk_level': self.calculate_cvss_risk(normalized_score),
        'risk_factors': risk_factors
    }

def calculate_overall_risk(self, scan_data, vulnerabilities, cves, web_vulns, sql_vulns):
    """Calculate overall system risk score"""
    total_risk_score = 0
    risk_breakdown = {
        'port_risks': [],
        'critical_count': 0,
        'high_count': 0,
        'medium_count': 0,
        'low_count': 0
    }
    
    # Calculate risk for each port
    for port in scan_data.get('ports', []):
        port_risk = self.calculate_port_risk(port, vulnerabilities, cves)
        risk_breakdown['port_risks'].append(port_risk)
        total_risk_score += port_risk['risk_score']
    
    # Add web vulnerability risks
    for web_vuln in web_vulns:
        severity = web_vuln.get('severity', 'MEDIUM')
        total_risk_score += self.risk_scores.get(severity, 4)
        risk_breakdown[f"{severity.lower()}_count"] += 1
    
    # SQL injection = automatic CRITICAL
    if sql_vulns:
        total_risk_score += len(sql_vulns) * 10
        risk_breakdown['critical_count'] += len(sql_vulns)
    
    # Calculate overall risk level
    num_ports = len(scan_data.get('ports', []))
    avg_risk_score = total_risk_score / max(num_ports, 1)
    
    overall_risk = self.calculate_cvss_risk(avg_risk_score)
    
    return {
        'total_risk_score': round(total_risk_score, 2),
        'average_risk_score': round(avg_risk_score, 2),
        'overall_risk_level': overall_risk,
        'total_vulnerabilities': len(vulnerabilities),
        'web_vulnerabilities': len(web_vulns),
        'sql_vulnerabilities': len(sql_vulns),
        'port_risk_breakdown': risk_breakdown,
        'recommendations': self.generate_recommendations(risk_breakdown)
    }

def generate_recommendations(self, risk_breakdown):
    """Generate prioritized security recommendations"""
    recommendations = []
    
    # Critical vulnerabilities
    if risk_breakdown['critical_count'] > 0:
        recommendations.append({
            'priority': 'CRITICAL',
            'action': 'Patch critical vulnerabilities immediately',
            'details': f"{risk_breakdown['critical_count']} critical issues found"
        })
    
    # Sensitive service exposure
    for port_risk in risk_breakdown['port_risks']:
        if any('Sensitive service' in factor for factor in port_risk['risk_factors']):
            recommendations.append({
                'priority': 'HIGH',
                'action': f"Secure or disable {port_risk['service']} on port {port_risk['port']}",
                'details': 'Consider firewall restrictions or VPN-only access'
            })
    
    # Exploitable vulnerabilities
    exploitable_ports = [p for p in risk_breakdown['port_risks'] 
                        if any('Exploitable' in f for f in p['risk_factors'])]
    if exploitable_ports:
        recommendations.append({
            'priority': 'HIGH',
            'action': 'Address exploitable vulnerabilities',
            'details': f"{len(exploitable_ports)} ports have exploitable issues"
        })
    
    return recommendations