1use axum::{
7 extract::{Path, Query, State},
8 http::StatusCode,
9 response::Json,
10};
11use mockforge_core::security::{
12 emit_security_event,
13 risk_assessment::{
14 Impact, Likelihood, Risk, RiskAssessmentEngine, RiskCategory, RiskLevel, TreatmentOption,
15 TreatmentStatus,
16 },
17 EventActor, EventOutcome, EventTarget, SecurityEvent, SecurityEventType,
18};
19use serde::{Deserialize, Serialize};
20use std::collections::HashMap;
21use std::sync::Arc;
22use tokio::sync::RwLock;
23use tracing::{error, info};
24
25use crate::handlers::auth_helpers::{extract_user_id_with_fallback, OptionalAuthClaims};
26
27#[derive(Clone)]
29pub struct RiskAssessmentState {
30 pub engine: Arc<RwLock<RiskAssessmentEngine>>,
32}
33
34#[derive(Debug, Deserialize)]
36pub struct CreateRiskRequest {
37 pub title: String,
39 pub description: String,
41 pub category: RiskCategory,
43 pub subcategory: Option<String>,
45 pub likelihood: Likelihood,
47 pub impact: Impact,
49 pub threat: Option<String>,
51 pub vulnerability: Option<String>,
53 pub asset: Option<String>,
55 pub existing_controls: Option<Vec<String>>,
57 pub compliance_requirements: Option<Vec<String>>,
59}
60
61#[derive(Debug, Deserialize)]
63pub struct UpdateRiskAssessmentRequest {
64 pub likelihood: Option<Likelihood>,
66 pub impact: Option<Impact>,
68}
69
70#[derive(Debug, Deserialize)]
72pub struct UpdateTreatmentPlanRequest {
73 pub treatment_option: TreatmentOption,
75 pub treatment_plan: Vec<String>,
77 pub treatment_owner: Option<String>,
79 pub treatment_deadline: Option<chrono::DateTime<chrono::Utc>>,
81}
82
83#[derive(Debug, Deserialize)]
85pub struct SetResidualRiskRequest {
86 pub residual_likelihood: Likelihood,
88 pub residual_impact: Impact,
90}
91
92#[derive(Debug, Serialize)]
94pub struct RiskListResponse {
95 pub risks: Vec<Risk>,
97 pub summary: mockforge_core::security::risk_assessment::RiskSummary,
99}
100
101pub async fn create_risk(
105 State(state): State<RiskAssessmentState>,
106 claims: OptionalAuthClaims,
107 Json(request): Json<CreateRiskRequest>,
108) -> Result<Json<serde_json::Value>, StatusCode> {
109 let created_by = extract_user_id_with_fallback(&claims);
111
112 let engine = state.engine.write().await;
113 let risk = engine
114 .create_risk(
115 request.title,
116 request.description,
117 request.category,
118 request.likelihood,
119 request.impact,
120 created_by,
121 )
122 .await
123 .map_err(|e| {
124 error!("Failed to create risk: {}", e);
125 StatusCode::INTERNAL_SERVER_ERROR
126 })?;
127
128 if let Some(subcategory) = request.subcategory {
130 }
133 if let Some(threat) = request.threat {
134 }
136 if let Some(vulnerability) = request.vulnerability {
137 }
139 if let Some(asset) = request.asset {
140 }
142 if let Some(controls) = request.existing_controls {
143 }
145 if let Some(requirements) = request.compliance_requirements {
146 }
148
149 info!("Risk created: {}", risk.risk_id);
150
151 let event = SecurityEvent::new(SecurityEventType::ConfigChanged, None, None)
153 .with_actor(EventActor {
154 user_id: Some(created_by.to_string()),
155 username: None,
156 ip_address: None,
157 user_agent: None,
158 })
159 .with_target(EventTarget {
160 resource_type: Some("risk".to_string()),
161 resource_id: Some(risk.risk_id.clone()),
162 method: None,
163 })
164 .with_outcome(EventOutcome {
165 success: true,
166 reason: Some("Risk created".to_string()),
167 });
168 emit_security_event(event).await;
169
170 Ok(Json(serde_json::to_value(&risk).map_err(|e| {
171 error!("Failed to serialize risk: {}", e);
172 StatusCode::INTERNAL_SERVER_ERROR
173 })?))
174}
175
176pub async fn get_risk(
180 State(state): State<RiskAssessmentState>,
181 Path(risk_id): Path<String>,
182) -> Result<Json<serde_json::Value>, StatusCode> {
183 let engine = state.engine.read().await;
184 let risk = engine
185 .get_risk(&risk_id)
186 .await
187 .map_err(|e| {
188 error!("Failed to get risk: {}", e);
189 StatusCode::INTERNAL_SERVER_ERROR
190 })?
191 .ok_or_else(|| {
192 error!("Risk not found: {}", risk_id);
193 StatusCode::NOT_FOUND
194 })?;
195
196 Ok(Json(serde_json::to_value(&risk).map_err(|e| {
197 error!("Failed to serialize risk: {}", e);
198 StatusCode::INTERNAL_SERVER_ERROR
199 })?))
200}
201
202pub async fn list_risks(
206 State(state): State<RiskAssessmentState>,
207 Query(params): Query<HashMap<String, String>>,
208) -> Result<Json<RiskListResponse>, StatusCode> {
209 let engine = state.engine.read().await;
210
211 let risks = if let Some(level_str) = params.get("level") {
212 let level = match level_str.as_str() {
213 "critical" => RiskLevel::Critical,
214 "high" => RiskLevel::High,
215 "medium" => RiskLevel::Medium,
216 "low" => RiskLevel::Low,
217 _ => return Err(StatusCode::BAD_REQUEST),
218 };
219 engine.get_risks_by_level(level).await.map_err(|e| {
220 error!("Failed to get risks by level: {}", e);
221 StatusCode::INTERNAL_SERVER_ERROR
222 })?
223 } else if let Some(category_str) = params.get("category") {
224 let category = match category_str.as_str() {
225 "technical" => RiskCategory::Technical,
226 "operational" => RiskCategory::Operational,
227 "compliance" => RiskCategory::Compliance,
228 "business" => RiskCategory::Business,
229 _ => return Err(StatusCode::BAD_REQUEST),
230 };
231 engine.get_risks_by_category(category).await.map_err(|e| {
232 error!("Failed to get risks by category: {}", e);
233 StatusCode::INTERNAL_SERVER_ERROR
234 })?
235 } else if let Some(status_str) = params.get("treatment_status") {
236 let status = match status_str.as_str() {
237 "not_started" => TreatmentStatus::NotStarted,
238 "in_progress" => TreatmentStatus::InProgress,
239 "completed" => TreatmentStatus::Completed,
240 "on_hold" => TreatmentStatus::OnHold,
241 _ => return Err(StatusCode::BAD_REQUEST),
242 };
243 engine.get_risks_by_treatment_status(status).await.map_err(|e| {
244 error!("Failed to get risks by treatment status: {}", e);
245 StatusCode::INTERNAL_SERVER_ERROR
246 })?
247 } else {
248 engine.get_all_risks().await.map_err(|e| {
249 error!("Failed to get all risks: {}", e);
250 StatusCode::INTERNAL_SERVER_ERROR
251 })?
252 };
253
254 let summary = engine.get_risk_summary().await.map_err(|e| {
255 error!("Failed to get risk summary: {}", e);
256 StatusCode::INTERNAL_SERVER_ERROR
257 })?;
258
259 Ok(Json(RiskListResponse { risks, summary }))
260}
261
262pub async fn update_risk_assessment(
266 State(state): State<RiskAssessmentState>,
267 Path(risk_id): Path<String>,
268 claims: OptionalAuthClaims,
269 Json(request): Json<UpdateRiskAssessmentRequest>,
270) -> Result<Json<serde_json::Value>, StatusCode> {
271 let updated_by = extract_user_id_with_fallback(&claims);
273
274 let engine = state.engine.write().await;
275 engine
276 .update_risk_assessment(&risk_id, request.likelihood, request.impact)
277 .await
278 .map_err(|e| {
279 error!("Failed to update risk assessment: {}", e);
280 StatusCode::BAD_REQUEST
281 })?;
282
283 info!("Risk assessment updated: {}", risk_id);
284
285 let event = SecurityEvent::new(SecurityEventType::ConfigChanged, None, None)
287 .with_actor(EventActor {
288 user_id: Some(updated_by.to_string()),
289 username: None,
290 ip_address: None,
291 user_agent: None,
292 })
293 .with_target(EventTarget {
294 resource_type: Some("risk".to_string()),
295 resource_id: Some(risk_id.clone()),
296 method: None,
297 })
298 .with_outcome(EventOutcome {
299 success: true,
300 reason: Some("Risk assessment updated".to_string()),
301 });
302 emit_security_event(event).await;
303
304 Ok(Json(serde_json::json!({
305 "risk_id": risk_id,
306 "status": "updated"
307 })))
308}
309
310pub async fn update_treatment_plan(
314 State(state): State<RiskAssessmentState>,
315 Path(risk_id): Path<String>,
316 claims: OptionalAuthClaims,
317 Json(request): Json<UpdateTreatmentPlanRequest>,
318) -> Result<Json<serde_json::Value>, StatusCode> {
319 let updated_by = extract_user_id_with_fallback(&claims);
321
322 let engine = state.engine.write().await;
323 engine
324 .update_treatment_plan(
325 &risk_id,
326 request.treatment_option,
327 request.treatment_plan,
328 request.treatment_owner,
329 request.treatment_deadline,
330 )
331 .await
332 .map_err(|e| {
333 error!("Failed to update treatment plan: {}", e);
334 StatusCode::BAD_REQUEST
335 })?;
336
337 info!("Treatment plan updated: {}", risk_id);
338
339 let event = SecurityEvent::new(SecurityEventType::ConfigChanged, None, None)
341 .with_actor(EventActor {
342 user_id: Some(updated_by.to_string()),
343 username: None,
344 ip_address: None,
345 user_agent: None,
346 })
347 .with_target(EventTarget {
348 resource_type: Some("risk".to_string()),
349 resource_id: Some(risk_id.clone()),
350 method: None,
351 })
352 .with_outcome(EventOutcome {
353 success: true,
354 reason: Some("Treatment plan updated".to_string()),
355 });
356 emit_security_event(event).await;
357
358 Ok(Json(serde_json::json!({
359 "risk_id": risk_id,
360 "status": "updated"
361 })))
362}
363
364pub async fn update_treatment_status(
368 State(state): State<RiskAssessmentState>,
369 Path(risk_id): Path<String>,
370 claims: OptionalAuthClaims,
371 Json(request): Json<serde_json::Value>,
372) -> Result<Json<serde_json::Value>, StatusCode> {
373 let _updated_by = extract_user_id_with_fallback(&claims);
375
376 let status_str =
377 request.get("status").and_then(|v| v.as_str()).ok_or(StatusCode::BAD_REQUEST)?;
378
379 let status = match status_str {
380 "not_started" => TreatmentStatus::NotStarted,
381 "in_progress" => TreatmentStatus::InProgress,
382 "completed" => TreatmentStatus::Completed,
383 "on_hold" => TreatmentStatus::OnHold,
384 _ => return Err(StatusCode::BAD_REQUEST),
385 };
386
387 let engine = state.engine.write().await;
388 engine.update_treatment_status(&risk_id, status).await.map_err(|e| {
389 error!("Failed to update treatment status: {}", e);
390 StatusCode::BAD_REQUEST
391 })?;
392
393 info!("Treatment status updated: {}", risk_id);
394
395 Ok(Json(serde_json::json!({
396 "risk_id": risk_id,
397 "status": "updated"
398 })))
399}
400
401pub async fn set_residual_risk(
405 State(state): State<RiskAssessmentState>,
406 Path(risk_id): Path<String>,
407 claims: OptionalAuthClaims,
408 Json(request): Json<SetResidualRiskRequest>,
409) -> Result<Json<serde_json::Value>, StatusCode> {
410 let _updated_by = extract_user_id_with_fallback(&claims);
412
413 let engine = state.engine.write().await;
414 engine
415 .set_residual_risk(&risk_id, request.residual_likelihood, request.residual_impact)
416 .await
417 .map_err(|e| {
418 error!("Failed to set residual risk: {}", e);
419 StatusCode::BAD_REQUEST
420 })?;
421
422 info!("Residual risk set: {}", risk_id);
423
424 Ok(Json(serde_json::json!({
425 "risk_id": risk_id,
426 "status": "updated"
427 })))
428}
429
430pub async fn review_risk(
434 State(state): State<RiskAssessmentState>,
435 Path(risk_id): Path<String>,
436 claims: OptionalAuthClaims,
437) -> Result<Json<serde_json::Value>, StatusCode> {
438 let reviewed_by = extract_user_id_with_fallback(&claims);
440
441 let engine = state.engine.write().await;
442 engine.review_risk(&risk_id, reviewed_by).await.map_err(|e| {
443 error!("Failed to review risk: {}", e);
444 StatusCode::BAD_REQUEST
445 })?;
446
447 info!("Risk reviewed: {}", risk_id);
448
449 Ok(Json(serde_json::json!({
450 "risk_id": risk_id,
451 "status": "reviewed"
452 })))
453}
454
455pub async fn get_risks_due_for_review(
459 State(state): State<RiskAssessmentState>,
460) -> Result<Json<serde_json::Value>, StatusCode> {
461 let engine = state.engine.read().await;
462 let risks = engine.get_risks_due_for_review().await.map_err(|e| {
463 error!("Failed to get risks due for review: {}", e);
464 StatusCode::INTERNAL_SERVER_ERROR
465 })?;
466
467 Ok(Json(serde_json::to_value(&risks).map_err(|e| {
468 error!("Failed to serialize risks: {}", e);
469 StatusCode::INTERNAL_SERVER_ERROR
470 })?))
471}
472
473pub async fn get_risk_summary(
477 State(state): State<RiskAssessmentState>,
478) -> Result<Json<serde_json::Value>, StatusCode> {
479 let engine = state.engine.read().await;
480 let summary = engine.get_risk_summary().await.map_err(|e| {
481 error!("Failed to get risk summary: {}", e);
482 StatusCode::INTERNAL_SERVER_ERROR
483 })?;
484
485 Ok(Json(serde_json::to_value(&summary).map_err(|e| {
486 error!("Failed to serialize summary: {}", e);
487 StatusCode::INTERNAL_SERVER_ERROR
488 })?))
489}
490
491pub fn risk_assessment_router(state: RiskAssessmentState) -> axum::Router {
493 use axum::routing::{get, patch, post, put};
494
495 axum::Router::new()
496 .route("/risks", get(list_risks))
497 .route("/risks", post(create_risk))
498 .route("/risks/{risk_id}", get(get_risk))
499 .route("/risks/{risk_id}/assessment", put(update_risk_assessment))
500 .route("/risks/{risk_id}/treatment", put(update_treatment_plan))
501 .route("/risks/{risk_id}/treatment/status", patch(update_treatment_status))
502 .route("/risks/{risk_id}/residual", put(set_residual_risk))
503 .route("/risks/{risk_id}/review", post(review_risk))
504 .route("/risks/due-for-review", get(get_risks_due_for_review))
505 .route("/risks/summary", get(get_risk_summary))
506 .with_state(state)
507}