1use crate::result::{DetectionResult, Severity};
4use std::fmt;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
11pub enum RiskLevel {
12 None,
13 Low,
14 Medium,
15 High,
16 Critical,
17}
18
19impl fmt::Display for RiskLevel {
20 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21 f.write_str(match self {
22 RiskLevel::None => "NONE",
23 RiskLevel::Low => "LOW",
24 RiskLevel::Medium => "MEDIUM",
25 RiskLevel::High => "HIGH",
26 RiskLevel::Critical => "CRITICAL",
27 })
28 }
29}
30
31fn weight(severity: &Severity) -> u32 {
34 match severity {
35 Severity::Critical => 100,
36 Severity::High => 40,
37 Severity::Medium => 15,
38 Severity::Low => 5,
39 }
40}
41
42#[must_use]
44pub fn total(results: &[DetectionResult]) -> u32 {
45 results.iter().map(|r| weight(&r.severity)).sum()
46}
47
48#[must_use]
54pub fn score(results: &[DetectionResult]) -> RiskLevel {
55 if results.is_empty() {
56 return RiskLevel::None;
57 }
58 if results.iter().any(|r| r.severity == Severity::Critical) {
59 return RiskLevel::Critical;
60 }
61 level_of(total(results))
62}
63
64fn level_of(points: u32) -> RiskLevel {
65 match points {
66 0 => RiskLevel::None,
67 1..=14 => RiskLevel::Low,
68 15..=39 => RiskLevel::Medium,
69 40..=99 => RiskLevel::High,
70 _ => RiskLevel::Critical,
71 }
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub struct RiskAssessment {
77 pub level: RiskLevel,
78 pub score: u32,
79 pub results: usize,
80}
81
82#[must_use]
84pub fn assess(results: &[DetectionResult]) -> RiskAssessment {
85 RiskAssessment {
86 level: score(results),
87 score: total(results),
88 results: results.len(),
89 }
90}
91
92#[cfg(test)]
93mod tests {
94 use super::*;
95 use crate::result::AttackCategory;
96
97 fn hit(severity: Severity) -> DetectionResult {
98 DetectionResult {
99 attack_type: "test".into(),
100 category: AttackCategory::Injection,
101 severity,
102 matched_pattern: "x".into(),
103 offset: 0,
104 message: "test".into(),
105 }
106 }
107
108 fn hits(severities: &[Severity]) -> Vec<DetectionResult> {
109 severities.iter().map(|s| hit(*s)).collect()
110 }
111
112 fn repeated(severity: Severity, n: usize) -> Vec<DetectionResult> {
113 (0..n).map(|_| hit(severity)).collect()
114 }
115
116 #[test]
117 fn empty_is_none() {
118 assert_eq!(score(&[]), RiskLevel::None);
119 assert_eq!(total(&[]), 0);
120 assert_eq!(
121 assess(&[]),
122 RiskAssessment {
123 level: RiskLevel::None,
124 score: 0,
125 results: 0
126 }
127 );
128 }
129
130 #[test]
131 fn single_critical_is_critical() {
132 assert_eq!(score(&hits(&[Severity::Critical])), RiskLevel::Critical);
133 }
134
135 #[test]
136 fn critical_short_circuits_even_with_low() {
137 assert_eq!(
138 score(&hits(&[Severity::Low, Severity::Critical, Severity::Low])),
139 RiskLevel::Critical
140 );
141 }
142
143 #[test]
144 fn single_low_does_not_escalate() {
145 assert_eq!(score(&hits(&[Severity::Low])), RiskLevel::Low);
146 assert_eq!(score(&hits(&[Severity::Medium])), RiskLevel::Medium);
147 assert_eq!(score(&hits(&[Severity::High])), RiskLevel::High);
148 }
149
150 #[test]
151 fn stacked_lows_escalate() {
152 assert_eq!(
154 score(&hits(&[Severity::Low, Severity::Low])),
155 RiskLevel::Low
156 );
157 assert_eq!(
159 score(&hits(&[Severity::Low, Severity::Low, Severity::Low])),
160 RiskLevel::Medium
161 );
162 assert_eq!(
164 score(&repeated(Severity::Low, 8)),
165 RiskLevel::High,
166 "8 × Low 应升级到 High"
167 );
168 }
169
170 #[test]
171 fn stacked_mediums_and_highs_escalate() {
172 assert_eq!(score(&repeated(Severity::Medium, 3)), RiskLevel::High);
174 assert_eq!(score(&repeated(Severity::High, 3)), RiskLevel::Critical);
176 assert_eq!(
178 score(&hits(&[Severity::High, Severity::Low])),
179 RiskLevel::High
180 );
181 }
182
183 #[test]
184 fn boundaries() {
185 assert_eq!(level_of(0), RiskLevel::None);
186 assert_eq!(level_of(1), RiskLevel::Low);
187 assert_eq!(level_of(14), RiskLevel::Low);
188 assert_eq!(level_of(15), RiskLevel::Medium);
189 assert_eq!(level_of(39), RiskLevel::Medium);
190 assert_eq!(level_of(40), RiskLevel::High);
191 assert_eq!(level_of(99), RiskLevel::High);
192 assert_eq!(level_of(100), RiskLevel::Critical);
193 }
194
195 #[test]
196 fn assessment_carries_counts_and_points() {
197 let a = assess(&hits(&[Severity::High, Severity::Low, Severity::Low]));
198 assert_eq!(a.results, 3);
199 assert_eq!(a.score, 50);
200 assert_eq!(a.level, RiskLevel::High);
201 }
202
203 #[test]
204 fn risk_level_ordering() {
205 assert!(RiskLevel::None < RiskLevel::Low);
206 assert!(RiskLevel::Low < RiskLevel::Medium);
207 assert!(RiskLevel::Medium < RiskLevel::High);
208 assert!(RiskLevel::High < RiskLevel::Critical);
209 }
210
211 #[test]
212 fn risk_level_display_uppercase() {
213 assert_eq!(RiskLevel::None.to_string(), "NONE");
214 assert_eq!(RiskLevel::Low.to_string(), "LOW");
215 assert_eq!(RiskLevel::Medium.to_string(), "MEDIUM");
216 assert_eq!(RiskLevel::High.to_string(), "HIGH");
217 assert_eq!(RiskLevel::Critical.to_string(), "CRITICAL");
218 }
219}