1use crate::Error;
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use uuid::Uuid;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum RiskCategory {
16 Technical,
18 Operational,
20 Compliance,
22 Business,
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
28#[serde(rename_all = "lowercase")]
29pub enum Likelihood {
30 Rare = 1,
32 Unlikely = 2,
34 Possible = 3,
36 Likely = 4,
38 AlmostCertain = 5,
40}
41
42impl Likelihood {
43 pub fn value(&self) -> u8 {
45 *self as u8
46 }
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
51#[serde(rename_all = "lowercase")]
52pub enum Impact {
53 Negligible = 1,
55 Low = 2,
57 Medium = 3,
59 High = 4,
61 Critical = 5,
63}
64
65impl Impact {
66 pub fn value(&self) -> u8 {
68 *self as u8
69 }
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
74#[serde(rename_all = "lowercase")]
75pub enum RiskLevel {
76 Low,
78 Medium,
80 High,
82 Critical,
84}
85
86impl RiskLevel {
87 pub fn from_score(score: u8) -> Self {
89 match score {
90 1..=5 => RiskLevel::Low,
91 6..=11 => RiskLevel::Medium,
92 12..=19 => RiskLevel::High,
93 20..=25 => RiskLevel::Critical,
94 _ => RiskLevel::Low,
95 }
96 }
97}
98
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
101#[serde(rename_all = "lowercase")]
102pub enum TreatmentOption {
103 Avoid,
105 Mitigate,
107 Transfer,
109 Accept,
111}
112
113#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
115#[serde(rename_all = "snake_case")]
116pub enum TreatmentStatus {
117 NotStarted,
119 InProgress,
121 Completed,
123 OnHold,
125}
126
127#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
129#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
130#[serde(rename_all = "lowercase")]
131pub enum RiskReviewFrequency {
132 Monthly,
134 Quarterly,
136 Annually,
138 AdHoc,
140}
141
142#[derive(Debug, Clone, Serialize, Deserialize)]
144pub struct Risk {
145 pub risk_id: String,
147 pub title: String,
149 pub description: String,
151 pub category: RiskCategory,
153 pub subcategory: Option<String>,
155 pub likelihood: Likelihood,
157 pub impact: Impact,
159 pub risk_score: u8,
161 pub risk_level: RiskLevel,
163 pub threat: Option<String>,
165 pub vulnerability: Option<String>,
167 pub asset: Option<String>,
169 pub existing_controls: Vec<String>,
171 pub treatment_option: TreatmentOption,
173 pub treatment_plan: Vec<String>,
175 pub treatment_owner: Option<String>,
177 pub treatment_deadline: Option<DateTime<Utc>>,
179 pub treatment_status: TreatmentStatus,
181 pub residual_likelihood: Option<Likelihood>,
183 pub residual_impact: Option<Impact>,
185 pub residual_risk_score: Option<u8>,
187 pub residual_risk_level: Option<RiskLevel>,
189 pub last_reviewed: Option<DateTime<Utc>>,
191 pub next_review: Option<DateTime<Utc>>,
193 pub review_frequency: RiskReviewFrequency,
195 pub compliance_requirements: Vec<String>,
197 pub created_at: DateTime<Utc>,
199 pub updated_at: DateTime<Utc>,
201 pub created_by: Uuid,
203}
204
205impl Risk {
206 pub fn new(
208 risk_id: String,
209 title: String,
210 description: String,
211 category: RiskCategory,
212 likelihood: Likelihood,
213 impact: Impact,
214 created_by: Uuid,
215 ) -> Self {
216 let risk_score = likelihood.value() * impact.value();
217 let risk_level = RiskLevel::from_score(risk_score);
218
219 Self {
220 risk_id,
221 title,
222 description,
223 category,
224 subcategory: None,
225 likelihood,
226 impact,
227 risk_score,
228 risk_level,
229 threat: None,
230 vulnerability: None,
231 asset: None,
232 existing_controls: Vec::new(),
233 treatment_option: TreatmentOption::Accept,
234 treatment_plan: Vec::new(),
235 treatment_owner: None,
236 treatment_deadline: None,
237 treatment_status: TreatmentStatus::NotStarted,
238 residual_likelihood: None,
239 residual_impact: None,
240 residual_risk_score: None,
241 residual_risk_level: None,
242 last_reviewed: None,
243 next_review: None,
244 review_frequency: RiskReviewFrequency::Quarterly,
245 compliance_requirements: Vec::new(),
246 created_at: Utc::now(),
247 updated_at: Utc::now(),
248 created_by,
249 }
250 }
251
252 pub fn recalculate(&mut self) {
254 self.risk_score = self.likelihood.value() * self.impact.value();
255 self.risk_level = RiskLevel::from_score(self.risk_score);
256
257 if let (Some(res_likelihood), Some(res_impact)) =
258 (self.residual_likelihood, self.residual_impact)
259 {
260 self.residual_risk_score = Some(res_likelihood.value() * res_impact.value());
261 self.residual_risk_level = self.residual_risk_score.map(RiskLevel::from_score);
262 }
263 }
264
265 pub fn calculate_next_review(&mut self) {
267 let now = Utc::now();
268 let next = match self.review_frequency {
269 RiskReviewFrequency::Monthly => now + chrono::Duration::days(30),
270 RiskReviewFrequency::Quarterly => now + chrono::Duration::days(90),
271 RiskReviewFrequency::Annually => now + chrono::Duration::days(365),
272 RiskReviewFrequency::AdHoc => now + chrono::Duration::days(90), };
274 self.next_review = Some(next);
275 }
276}
277
278#[derive(Debug, Clone, Serialize, Deserialize)]
280pub struct RiskSummary {
281 pub total_risks: u32,
283 pub critical: u32,
285 pub high: u32,
287 pub medium: u32,
289 pub low: u32,
291 pub by_category: HashMap<RiskCategory, u32>,
293 pub by_treatment_status: HashMap<TreatmentStatus, u32>,
295}
296
297#[derive(Debug, Clone, Serialize, Deserialize)]
299#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
300pub struct RiskAssessmentConfig {
301 pub enabled: bool,
303 pub default_review_frequency: RiskReviewFrequency,
305 pub risk_tolerance: RiskTolerance,
307}
308
309#[derive(Debug, Clone, Serialize, Deserialize)]
311#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
312pub struct RiskTolerance {
313 pub max_acceptable_score: u8,
315 pub require_treatment_above: u8,
317}
318
319impl Default for RiskAssessmentConfig {
320 fn default() -> Self {
321 Self {
322 enabled: true,
323 default_review_frequency: RiskReviewFrequency::Quarterly,
324 risk_tolerance: RiskTolerance {
325 max_acceptable_score: 5, require_treatment_above: 11, },
328 }
329 }
330}
331
332pub struct RiskAssessmentEngine {
334 config: RiskAssessmentConfig,
335 risks: std::sync::Arc<tokio::sync::RwLock<HashMap<String, Risk>>>,
337 risk_id_counter: std::sync::Arc<tokio::sync::RwLock<u64>>,
339}
340
341impl RiskAssessmentEngine {
342 pub fn new(config: RiskAssessmentConfig) -> Self {
344 Self {
345 config,
346 risks: std::sync::Arc::new(tokio::sync::RwLock::new(HashMap::new())),
347 risk_id_counter: std::sync::Arc::new(tokio::sync::RwLock::new(0)),
348 }
349 }
350
351 async fn generate_risk_id(&self) -> String {
353 let mut counter = self.risk_id_counter.write().await;
354 *counter += 1;
355 format!("RISK-{:03}", *counter)
356 }
357
358 pub async fn create_risk(
360 &self,
361 title: String,
362 description: String,
363 category: RiskCategory,
364 likelihood: Likelihood,
365 impact: Impact,
366 created_by: Uuid,
367 ) -> Result<Risk, Error> {
368 let risk_id = self.generate_risk_id().await;
369 let mut risk = Risk::new(
370 risk_id.clone(),
371 title,
372 description,
373 category,
374 likelihood,
375 impact,
376 created_by,
377 );
378 risk.review_frequency = self.config.default_review_frequency;
379 risk.calculate_next_review();
380
381 let mut risks = self.risks.write().await;
382 risks.insert(risk_id, risk.clone());
383
384 Ok(risk)
385 }
386
387 pub async fn get_risk(&self, risk_id: &str) -> Result<Option<Risk>, Error> {
389 let risks = self.risks.read().await;
390 Ok(risks.get(risk_id).cloned())
391 }
392
393 pub async fn get_all_risks(&self) -> Result<Vec<Risk>, Error> {
395 let risks = self.risks.read().await;
396 Ok(risks.values().cloned().collect())
397 }
398
399 pub async fn get_risks_by_level(&self, level: RiskLevel) -> Result<Vec<Risk>, Error> {
401 let risks = self.risks.read().await;
402 Ok(risks.values().filter(|r| r.risk_level == level).cloned().collect())
403 }
404
405 pub async fn get_risks_by_category(&self, category: RiskCategory) -> Result<Vec<Risk>, Error> {
407 let risks = self.risks.read().await;
408 Ok(risks.values().filter(|r| r.category == category).cloned().collect())
409 }
410
411 pub async fn get_risks_by_treatment_status(
413 &self,
414 status: TreatmentStatus,
415 ) -> Result<Vec<Risk>, Error> {
416 let risks = self.risks.read().await;
417 Ok(risks.values().filter(|r| r.treatment_status == status).cloned().collect())
418 }
419
420 pub async fn update_risk(&self, risk_id: &str, mut risk: Risk) -> Result<(), Error> {
422 risk.recalculate();
423 risk.updated_at = Utc::now();
424
425 let mut risks = self.risks.write().await;
426 if risks.contains_key(risk_id) {
427 risks.insert(risk_id.to_string(), risk);
428 Ok(())
429 } else {
430 Err(Error::Generic("Risk not found".to_string()))
431 }
432 }
433
434 pub async fn update_risk_assessment(
436 &self,
437 risk_id: &str,
438 likelihood: Option<Likelihood>,
439 impact: Option<Impact>,
440 ) -> Result<(), Error> {
441 let mut risks = self.risks.write().await;
442 if let Some(risk) = risks.get_mut(risk_id) {
443 if let Some(l) = likelihood {
444 risk.likelihood = l;
445 }
446 if let Some(i) = impact {
447 risk.impact = i;
448 }
449 risk.recalculate();
450 risk.updated_at = Utc::now();
451 Ok(())
452 } else {
453 Err(Error::Generic("Risk not found".to_string()))
454 }
455 }
456
457 pub async fn update_treatment_plan(
459 &self,
460 risk_id: &str,
461 treatment_option: TreatmentOption,
462 treatment_plan: Vec<String>,
463 treatment_owner: Option<String>,
464 treatment_deadline: Option<DateTime<Utc>>,
465 ) -> Result<(), Error> {
466 let mut risks = self.risks.write().await;
467 if let Some(risk) = risks.get_mut(risk_id) {
468 risk.treatment_option = treatment_option;
469 risk.treatment_plan = treatment_plan;
470 risk.treatment_owner = treatment_owner;
471 risk.treatment_deadline = treatment_deadline;
472 risk.updated_at = Utc::now();
473 Ok(())
474 } else {
475 Err(Error::Generic("Risk not found".to_string()))
476 }
477 }
478
479 pub async fn update_treatment_status(
481 &self,
482 risk_id: &str,
483 status: TreatmentStatus,
484 ) -> Result<(), Error> {
485 let mut risks = self.risks.write().await;
486 if let Some(risk) = risks.get_mut(risk_id) {
487 risk.treatment_status = status;
488 risk.updated_at = Utc::now();
489 Ok(())
490 } else {
491 Err(Error::Generic("Risk not found".to_string()))
492 }
493 }
494
495 pub async fn set_residual_risk(
497 &self,
498 risk_id: &str,
499 residual_likelihood: Likelihood,
500 residual_impact: Impact,
501 ) -> Result<(), Error> {
502 let mut risks = self.risks.write().await;
503 if let Some(risk) = risks.get_mut(risk_id) {
504 risk.residual_likelihood = Some(residual_likelihood);
505 risk.residual_impact = Some(residual_impact);
506 risk.recalculate();
507 risk.updated_at = Utc::now();
508 Ok(())
509 } else {
510 Err(Error::Generic("Risk not found".to_string()))
511 }
512 }
513
514 pub async fn review_risk(&self, risk_id: &str, reviewed_by: Uuid) -> Result<(), Error> {
516 let mut risks = self.risks.write().await;
517 if let Some(risk) = risks.get_mut(risk_id) {
518 risk.last_reviewed = Some(Utc::now());
519 risk.calculate_next_review();
520 risk.updated_at = Utc::now();
521 let _ = reviewed_by; Ok(())
523 } else {
524 Err(Error::Generic("Risk not found".to_string()))
525 }
526 }
527
528 pub async fn get_risk_summary(&self) -> Result<RiskSummary, Error> {
530 let risks = self.risks.read().await;
531
532 let mut summary = RiskSummary {
533 total_risks: risks.len() as u32,
534 critical: 0,
535 high: 0,
536 medium: 0,
537 low: 0,
538 by_category: HashMap::new(),
539 by_treatment_status: HashMap::new(),
540 };
541
542 for risk in risks.values() {
543 match risk.risk_level {
544 RiskLevel::Critical => summary.critical += 1,
545 RiskLevel::High => summary.high += 1,
546 RiskLevel::Medium => summary.medium += 1,
547 RiskLevel::Low => summary.low += 1,
548 }
549
550 *summary.by_category.entry(risk.category).or_insert(0) += 1;
551 let count = summary.by_treatment_status.entry(risk.treatment_status).or_insert(0);
552 *count += 1;
553 }
554
555 Ok(summary)
556 }
557
558 pub async fn get_risks_due_for_review(&self) -> Result<Vec<Risk>, Error> {
560 let risks = self.risks.read().await;
561 let now = Utc::now();
562
563 Ok(risks
564 .values()
565 .filter(|r| r.next_review.map(|next| next <= now).unwrap_or(false))
566 .cloned()
567 .collect())
568 }
569}
570
571#[cfg(test)]
572mod tests {
573 use super::*;
574
575 #[tokio::test]
576 async fn test_risk_creation() {
577 let config = RiskAssessmentConfig::default();
578 let engine = RiskAssessmentEngine::new(config);
579
580 let risk = engine
581 .create_risk(
582 "Test Risk".to_string(),
583 "Test description".to_string(),
584 RiskCategory::Technical,
585 Likelihood::Possible,
586 Impact::High,
587 Uuid::new_v4(),
588 )
589 .await
590 .unwrap();
591
592 assert_eq!(risk.risk_score, 12); assert_eq!(risk.risk_level, RiskLevel::High);
594 }
595
596 #[test]
597 fn test_risk_level_calculation() {
598 assert_eq!(RiskLevel::from_score(3), RiskLevel::Low);
599 assert_eq!(RiskLevel::from_score(9), RiskLevel::Medium);
600 assert_eq!(RiskLevel::from_score(15), RiskLevel::High);
601 assert_eq!(RiskLevel::from_score(22), RiskLevel::Critical);
602 }
603}