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
42pub fn total(results: &[DetectionResult]) -> u32 {
44 results.iter().map(|r| weight(&r.severity)).sum()
45}
46
47pub fn score(results: &[DetectionResult]) -> RiskLevel {
53 if results.is_empty() {
54 return RiskLevel::None;
55 }
56 if results.iter().any(|r| r.severity == Severity::Critical) {
57 return RiskLevel::Critical;
58 }
59 level_of(total(results))
60}
61
62fn level_of(points: u32) -> RiskLevel {
63 match points {
64 0 => RiskLevel::None,
65 1..=14 => RiskLevel::Low,
66 15..=39 => RiskLevel::Medium,
67 40..=99 => RiskLevel::High,
68 _ => RiskLevel::Critical,
69 }
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub struct RiskAssessment {
75 pub level: RiskLevel,
76 pub score: u32,
77 pub results: usize,
78}
79
80pub fn assess(results: &[DetectionResult]) -> RiskAssessment {
82 RiskAssessment {
83 level: score(results),
84 score: total(results),
85 results: results.len(),
86 }
87}
88
89#[cfg(test)]
90mod tests {
91 use super::*;
92 use crate::result::AttackCategory;
93
94 fn hit(severity: Severity) -> DetectionResult {
95 DetectionResult {
96 attack_type: "test".into(),
97 category: AttackCategory::Injection,
98 severity,
99 matched_pattern: "x".into(),
100 offset: 0,
101 message: "test".into(),
102 }
103 }
104
105 fn hits(severities: &[Severity]) -> Vec<DetectionResult> {
106 severities.iter().map(|s| hit(s.clone())).collect()
107 }
108
109 fn repeated(severity: Severity, n: usize) -> Vec<DetectionResult> {
110 (0..n).map(|_| hit(severity.clone())).collect()
111 }
112
113 #[test]
114 fn empty_is_none() {
115 assert_eq!(score(&[]), RiskLevel::None);
116 assert_eq!(total(&[]), 0);
117 assert_eq!(
118 assess(&[]),
119 RiskAssessment {
120 level: RiskLevel::None,
121 score: 0,
122 results: 0
123 }
124 );
125 }
126
127 #[test]
128 fn single_critical_is_critical() {
129 assert_eq!(score(&hits(&[Severity::Critical])), RiskLevel::Critical);
130 }
131
132 #[test]
133 fn critical_short_circuits_even_with_low() {
134 assert_eq!(
135 score(&hits(&[Severity::Low, Severity::Critical, Severity::Low])),
136 RiskLevel::Critical
137 );
138 }
139
140 #[test]
141 fn single_low_does_not_escalate() {
142 assert_eq!(score(&hits(&[Severity::Low])), RiskLevel::Low);
143 assert_eq!(score(&hits(&[Severity::Medium])), RiskLevel::Medium);
144 assert_eq!(score(&hits(&[Severity::High])), RiskLevel::High);
145 }
146
147 #[test]
148 fn stacked_lows_escalate() {
149 assert_eq!(
151 score(&hits(&[Severity::Low, Severity::Low])),
152 RiskLevel::Low
153 );
154 assert_eq!(
156 score(&hits(&[Severity::Low, Severity::Low, Severity::Low])),
157 RiskLevel::Medium
158 );
159 assert_eq!(
161 score(&repeated(Severity::Low, 8)),
162 RiskLevel::High,
163 "8 × Low 应升级到 High"
164 );
165 }
166
167 #[test]
168 fn stacked_mediums_and_highs_escalate() {
169 assert_eq!(score(&repeated(Severity::Medium, 3)), RiskLevel::High);
171 assert_eq!(score(&repeated(Severity::High, 3)), RiskLevel::Critical);
173 assert_eq!(
175 score(&hits(&[Severity::High, Severity::Low])),
176 RiskLevel::High
177 );
178 }
179
180 #[test]
181 fn boundaries() {
182 assert_eq!(level_of(0), RiskLevel::None);
183 assert_eq!(level_of(1), RiskLevel::Low);
184 assert_eq!(level_of(14), RiskLevel::Low);
185 assert_eq!(level_of(15), RiskLevel::Medium);
186 assert_eq!(level_of(39), RiskLevel::Medium);
187 assert_eq!(level_of(40), RiskLevel::High);
188 assert_eq!(level_of(99), RiskLevel::High);
189 assert_eq!(level_of(100), RiskLevel::Critical);
190 }
191
192 #[test]
193 fn assessment_carries_counts_and_points() {
194 let a = assess(&hits(&[Severity::High, Severity::Low, Severity::Low]));
195 assert_eq!(a.results, 3);
196 assert_eq!(a.score, 50);
197 assert_eq!(a.level, RiskLevel::High);
198 }
199
200 #[test]
201 fn risk_level_ordering() {
202 assert!(RiskLevel::None < RiskLevel::Low);
203 assert!(RiskLevel::Low < RiskLevel::Medium);
204 assert!(RiskLevel::Medium < RiskLevel::High);
205 assert!(RiskLevel::High < RiskLevel::Critical);
206 }
207
208 #[test]
209 fn risk_level_display_uppercase() {
210 assert_eq!(RiskLevel::None.to_string(), "NONE");
211 assert_eq!(RiskLevel::Low.to_string(), "LOW");
212 assert_eq!(RiskLevel::Medium.to_string(), "MEDIUM");
213 assert_eq!(RiskLevel::High.to_string(), "HIGH");
214 assert_eq!(RiskLevel::Critical.to_string(), "CRITICAL");
215 }
216}