1use crate::rules::{Category, Finding, Severity};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5const CRITICAL_WEIGHT: u32 = 40;
7const HIGH_WEIGHT: u32 = 20;
8const MEDIUM_WEIGHT: u32 = 10;
9const LOW_WEIGHT: u32 = 5;
10const MAX_SCORE: u32 = 100;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "lowercase")]
15pub enum RiskLevel {
16 Safe,
17 Low,
18 Medium,
19 High,
20 Critical,
21}
22
23impl RiskLevel {
24 pub fn from_score(score: u32) -> Self {
25 match score {
26 0 => RiskLevel::Safe,
27 1..=25 => RiskLevel::Low,
28 26..=50 => RiskLevel::Medium,
29 51..=75 => RiskLevel::High,
30 _ => RiskLevel::Critical,
31 }
32 }
33
34 pub fn as_str(&self) -> &'static str {
35 match self {
36 RiskLevel::Safe => "SAFE",
37 RiskLevel::Low => "LOW",
38 RiskLevel::Medium => "MEDIUM",
39 RiskLevel::High => "HIGH",
40 RiskLevel::Critical => "CRITICAL",
41 }
42 }
43}
44
45impl std::fmt::Display for RiskLevel {
46 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47 write!(f, "{}", self.as_str())
48 }
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct CategoryScore {
54 pub category: String,
55 pub score: u32,
56 pub findings_count: usize,
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct RiskScore {
62 pub total: u32,
64 pub level: RiskLevel,
66 pub by_category: Vec<CategoryScore>,
68 pub by_severity: SeverityBreakdown,
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct SeverityBreakdown {
74 pub critical: u32,
75 pub high: u32,
76 pub medium: u32,
77 pub low: u32,
78}
79
80impl RiskScore {
81 pub fn from_findings(findings: &[Finding]) -> Self {
83 let mut category_scores: HashMap<Category, (u32, usize)> = HashMap::new();
84 let mut severity_scores = SeverityBreakdown {
85 critical: 0,
86 high: 0,
87 medium: 0,
88 low: 0,
89 };
90
91 for finding in findings {
92 let weight = match finding.severity {
93 Severity::Critical => {
94 severity_scores.critical += CRITICAL_WEIGHT;
95 CRITICAL_WEIGHT
96 }
97 Severity::High => {
98 severity_scores.high += HIGH_WEIGHT;
99 HIGH_WEIGHT
100 }
101 Severity::Medium => {
102 severity_scores.medium += MEDIUM_WEIGHT;
103 MEDIUM_WEIGHT
104 }
105 Severity::Low => {
106 severity_scores.low += LOW_WEIGHT;
107 LOW_WEIGHT
108 }
109 };
110
111 let entry = category_scores.entry(finding.category).or_insert((0, 0));
112 entry.0 += weight;
113 entry.1 += 1;
114 }
115
116 let raw_total: u32 = category_scores.values().map(|(s, _)| *s).sum();
118
119 let total = raw_total.min(MAX_SCORE);
121
122 let mut by_category: Vec<CategoryScore> = category_scores
124 .into_iter()
125 .map(|(cat, (score, count))| CategoryScore {
126 category: cat.as_str().to_string(),
127 score: score.min(MAX_SCORE),
128 findings_count: count,
129 })
130 .collect();
131
132 by_category.sort_by(|a, b| b.score.cmp(&a.score));
134
135 RiskScore {
136 total,
137 level: RiskLevel::from_score(total),
138 by_category,
139 by_severity: severity_scores,
140 }
141 }
142
143 pub fn score_bar(&self, score: u32, max: u32) -> String {
145 let filled = ((score as f32 / max as f32) * 10.0).round() as usize;
146 let filled = filled.min(10);
147 let empty = 10 - filled;
148 format!("{}{}", "█".repeat(filled), "░".repeat(empty))
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155 use crate::rules::{Confidence, Location};
156
157 fn create_test_finding(severity: Severity, category: Category) -> Finding {
158 Finding {
159 id: "TEST-001".to_string(),
160 severity,
161 category,
162 confidence: Confidence::Firm,
163 name: "Test".to_string(),
164 location: Location {
165 file: "test.sh".to_string(),
166 line: 1,
167 column: None,
168 },
169 code: "test".to_string(),
170 message: "test".to_string(),
171 recommendation: "test".to_string(),
172 fix_hint: None,
173 cwe_ids: vec![],
174 rule_severity: None,
175 client: None,
176 context: None,
177 }
178 }
179
180 #[test]
181 fn test_empty_findings_safe() {
182 let score = RiskScore::from_findings(&[]);
183 assert_eq!(score.total, 0);
184 assert_eq!(score.level, RiskLevel::Safe);
185 }
186
187 #[test]
188 fn test_single_critical_finding() {
189 let findings = vec![create_test_finding(
190 Severity::Critical,
191 Category::Exfiltration,
192 )];
193 let score = RiskScore::from_findings(&findings);
194 assert_eq!(score.total, 40);
195 assert_eq!(score.level, RiskLevel::Medium);
196 }
197
198 #[test]
199 fn test_multiple_findings_caps_at_100() {
200 let findings = vec![
201 create_test_finding(Severity::Critical, Category::Exfiltration),
202 create_test_finding(Severity::Critical, Category::PrivilegeEscalation),
203 create_test_finding(Severity::Critical, Category::Persistence),
204 ];
205 let score = RiskScore::from_findings(&findings);
206 assert_eq!(score.total, 100);
207 assert_eq!(score.level, RiskLevel::Critical);
208 }
209
210 #[test]
211 fn test_risk_level_boundaries() {
212 assert_eq!(RiskLevel::from_score(0), RiskLevel::Safe);
213 assert_eq!(RiskLevel::from_score(1), RiskLevel::Low);
214 assert_eq!(RiskLevel::from_score(25), RiskLevel::Low);
215 assert_eq!(RiskLevel::from_score(26), RiskLevel::Medium);
216 assert_eq!(RiskLevel::from_score(50), RiskLevel::Medium);
217 assert_eq!(RiskLevel::from_score(51), RiskLevel::High);
218 assert_eq!(RiskLevel::from_score(75), RiskLevel::High);
219 assert_eq!(RiskLevel::from_score(76), RiskLevel::Critical);
220 assert_eq!(RiskLevel::from_score(100), RiskLevel::Critical);
221 }
222
223 #[test]
224 fn test_category_breakdown() {
225 let findings = vec![
226 create_test_finding(Severity::Critical, Category::Exfiltration),
227 create_test_finding(Severity::High, Category::Exfiltration),
228 create_test_finding(Severity::Medium, Category::Persistence),
229 ];
230 let score = RiskScore::from_findings(&findings);
231
232 assert_eq!(score.by_category.len(), 2);
233 assert_eq!(score.by_category[0].category, "exfiltration");
235 assert_eq!(score.by_category[0].score, 60); assert_eq!(score.by_category[0].findings_count, 2);
237 }
238
239 #[test]
240 fn test_severity_breakdown() {
241 let findings = vec![
242 create_test_finding(Severity::Critical, Category::Exfiltration),
243 create_test_finding(Severity::High, Category::PrivilegeEscalation),
244 create_test_finding(Severity::Medium, Category::Persistence),
245 create_test_finding(Severity::Low, Category::Overpermission),
246 ];
247 let score = RiskScore::from_findings(&findings);
248
249 assert_eq!(score.by_severity.critical, 40);
250 assert_eq!(score.by_severity.high, 20);
251 assert_eq!(score.by_severity.medium, 10);
252 assert_eq!(score.by_severity.low, 5);
253 }
254
255 #[test]
256 fn test_score_bar() {
257 let score = RiskScore::from_findings(&[]);
258 assert_eq!(score.score_bar(0, 100), "░░░░░░░░░░");
259 assert_eq!(score.score_bar(50, 100), "█████░░░░░");
260 assert_eq!(score.score_bar(100, 100), "██████████");
261 assert_eq!(score.score_bar(75, 100), "████████░░");
262 }
263
264 #[test]
265 fn test_risk_level_display() {
266 assert_eq!(format!("{}", RiskLevel::Safe), "SAFE");
267 assert_eq!(format!("{}", RiskLevel::Critical), "CRITICAL");
268 }
269}