kindly_guard_server/shield/
mod.rs

1// Copyright 2025 Kindly Software Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! Shield display module for security status visualization
15
16use anyhow::Result;
17use std::collections::VecDeque;
18use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
19use std::sync::{Arc, Mutex, Weak};
20use std::time::{Duration, Instant};
21use tracing::error;
22
23use crate::config::ShieldConfig;
24use crate::scanner::Threat;
25use crate::traits::SecurityEventProcessor;
26
27pub mod cli;
28pub mod display;
29pub mod universal_display;
30
31pub use cli::{CliShield, DisplayFormat, ShieldStatus};
32pub use display::ShieldDisplay;
33pub use universal_display::{UniversalDisplay, UniversalDisplayConfig, UniversalShieldStatus};
34
35/// Security shield that tracks protection status
36pub struct Shield {
37    active: AtomicBool,
38    start_time: Instant,
39    threats_blocked: AtomicU64,
40    recent_threats: Arc<Mutex<VecDeque<TimestampedThreat>>>,
41    config: ShieldConfig,
42    /// Whether advanced protection (event processor) is enabled
43    event_processor_enabled: AtomicBool,
44    /// Weak reference to event processor for correlation data
45    event_processor: Mutex<Option<Weak<dyn SecurityEventProcessor>>>,
46}
47
48/// Threat with timestamp
49#[derive(Clone)]
50struct TimestampedThreat {
51    threat: Threat,
52    timestamp: Instant,
53}
54
55/// Shield information snapshot
56#[derive(serde::Serialize, serde::Deserialize)]
57pub struct ShieldInfo {
58    pub active: bool,
59    #[serde(with = "serde_duration")]
60    pub uptime: Duration,
61    pub threats_blocked: u64,
62    pub recent_threat_rate: f64,
63}
64
65// Custom serialization for Duration
66mod serde_duration {
67    use serde::{Deserialize, Deserializer, Serializer};
68    use std::time::Duration;
69
70    pub fn serialize<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
71    where
72        S: Serializer,
73    {
74        serializer.serialize_u64(duration.as_secs())
75    }
76
77    pub fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
78    where
79        D: Deserializer<'de>,
80    {
81        let secs = u64::deserialize(deserializer)?;
82        Ok(Duration::from_secs(secs))
83    }
84}
85
86/// Shield statistics
87pub struct ShieldStats {
88    pub threats_blocked: u64,
89    pub active: bool,
90}
91
92impl Default for Shield {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98impl Shield {
99    /// Create a new shield
100    pub fn new() -> Self {
101        Self::with_config(ShieldConfig::default())
102    }
103
104    /// Create a new shield with specific config
105    pub fn with_config(config: ShieldConfig) -> Self {
106        Self {
107            active: AtomicBool::new(false),
108            start_time: Instant::now(),
109            threats_blocked: AtomicU64::new(0),
110            recent_threats: Arc::new(Mutex::new(VecDeque::with_capacity(1000))),
111            config,
112            event_processor_enabled: AtomicBool::new(false),
113            event_processor: Mutex::new(None),
114        }
115    }
116
117    /// Set shield active status
118    pub fn set_active(&self, active: bool) {
119        self.active.store(active, Ordering::Relaxed);
120    }
121
122    /// Check if shield is active
123    pub fn is_active(&self) -> bool {
124        self.active.load(Ordering::Relaxed)
125    }
126
127    /// Set whether event processor (advanced protection) is enabled
128    pub fn set_event_processor_enabled(&self, enabled: bool) {
129        self.event_processor_enabled
130            .store(enabled, Ordering::Relaxed);
131    }
132
133    /// Check if event processor (advanced protection) is enabled
134    pub fn is_event_processor_enabled(&self) -> bool {
135        self.event_processor_enabled.load(Ordering::Relaxed)
136    }
137
138    /// Set event processor reference for correlation data
139    pub fn set_event_processor(&self, processor: &Arc<dyn SecurityEventProcessor>) {
140        match self.event_processor.lock() {
141            Ok(mut ep) => *ep = Some(Arc::downgrade(processor)),
142            Err(e) => error!("Failed to acquire event processor lock: {}", e),
143        }
144    }
145
146    /// Record detected threats
147    pub fn record_threats(&self, threats: &[Threat]) {
148        if threats.is_empty() {
149            return;
150        }
151
152        let count = threats.len() as u64;
153        self.threats_blocked.fetch_add(count, Ordering::Relaxed);
154
155        let now = Instant::now();
156        let Ok(mut recent) = self.recent_threats.lock() else {
157            error!("Failed to acquire recent threats lock");
158            return;
159        };
160
161        for threat in threats {
162            // Add to recent threats
163            recent.push_back(TimestampedThreat {
164                threat: threat.clone(),
165                timestamp: now,
166            });
167
168            // Keep only last N threats
169            while recent.len() > 1000 {
170                recent.pop_front();
171            }
172        }
173    }
174
175    /// Get shield information
176    pub fn get_info(&self) -> ShieldInfo {
177        let now = Instant::now();
178        let uptime = now.duration_since(self.start_time);
179
180        // Calculate recent threat rate (threats per minute in last 5 minutes)
181        let recent_rate = match self.recent_threats.lock() {
182            Ok(recent) => {
183                let five_mins_ago = now.checked_sub(Duration::from_secs(300)).unwrap_or(now); // If subtraction fails (shouldn't happen), use current time
184                let recent_count = recent
185                    .iter()
186                    .filter(|t| t.timestamp > five_mins_ago)
187                    .count() as f64;
188                recent_count / 5.0 // per minute
189            },
190            Err(e) => {
191                error!("Failed to acquire recent threats lock: {}", e);
192                0.0 // Default to 0 on error
193            },
194        };
195
196        // Check for attack patterns if processor is available
197        if self.is_event_processor_enabled() {
198            if let Ok(ep_lock) = self.event_processor.lock() {
199                if let Some(weak_proc) = ep_lock.as_ref() {
200                    if let Some(processor) = weak_proc.upgrade() {
201                        // Check if any client is under monitoring (attack detected)
202                        if processor.is_monitored("any") {
203                            tracing::trace!("Attack pattern correlation active");
204                        }
205                    }
206                }
207            }
208        }
209
210        ShieldInfo {
211            active: self.is_active(),
212            uptime,
213            threats_blocked: self.threats_blocked.load(Ordering::Relaxed),
214            recent_threat_rate: recent_rate,
215        }
216    }
217
218    /// Get recent threats
219    pub fn get_recent_threats(&self, limit: usize) -> Vec<Threat> {
220        match self.recent_threats.lock() {
221            Ok(recent) => recent
222                .iter()
223                .rev()
224                .take(limit)
225                .map(|t| t.threat.clone())
226                .collect(),
227            Err(e) => {
228                error!("Failed to acquire recent threats lock: {}", e);
229                vec![]
230            },
231        }
232    }
233
234    /// Get threat statistics by type
235    pub fn get_threat_stats(&self) -> std::collections::HashMap<crate::scanner::ThreatType, u64> {
236        use std::collections::HashMap;
237
238        match self.recent_threats.lock() {
239            Ok(recent) => {
240                let mut stats = HashMap::new();
241                for item in recent.iter() {
242                    *stats.entry(item.threat.threat_type.clone()).or_insert(0) += 1;
243                }
244                stats
245            },
246            Err(e) => {
247                error!("Failed to acquire recent threats lock: {}", e);
248                HashMap::new()
249            },
250        }
251    }
252
253    /// Start the shield display if configured
254    pub async fn start_display(self: Arc<Self>) -> Result<()> {
255        if !self.config.enabled {
256            return Ok(());
257        }
258
259        let display = ShieldDisplay::new(self.clone(), self.config.clone());
260        display.run().await
261    }
262
263    /// Get the start time of the shield
264    pub const fn start_time(&self) -> Instant {
265        self.start_time
266    }
267
268    /// Get shield statistics
269    pub fn stats(&self) -> ShieldStats {
270        ShieldStats {
271            threats_blocked: self.threats_blocked.load(Ordering::Relaxed),
272            active: self.is_active(),
273        }
274    }
275
276    /// Get the last threat type if any
277    pub fn last_threat_type(&self) -> Option<String> {
278        match self.recent_threats.lock() {
279            Ok(recent) => recent.back().map(|t| format!("{}", t.threat.threat_type)),
280            Err(e) => {
281                error!("Failed to acquire recent threats lock: {}", e);
282                None
283            },
284        }
285    }
286
287    /// Get scanner statistics
288    pub fn scanner_stats(&self) -> crate::scanner::ScannerStats {
289        // Return the threat counts tracked by the shield
290        match self.recent_threats.lock() {
291            Ok(threats) => {
292                let (unicode_count, injection_count) =
293                    threats
294                        .iter()
295                        .fold((0u64, 0u64), |(unicode, injection), item| {
296                            match &item.threat.threat_type {
297                                crate::scanner::ThreatType::UnicodeInvisible
298                                | crate::scanner::ThreatType::UnicodeBiDi
299                                | crate::scanner::ThreatType::UnicodeHomograph => {
300                                    (unicode + 1, injection)
301                                },
302
303                                crate::scanner::ThreatType::SqlInjection
304                                | crate::scanner::ThreatType::CommandInjection
305                                | crate::scanner::ThreatType::PromptInjection
306                                | crate::scanner::ThreatType::PathTraversal => {
307                                    (unicode, injection + 1)
308                                },
309
310                                _ => (unicode, injection),
311                            }
312                        });
313
314                crate::scanner::ScannerStats {
315                    unicode_threats_detected: unicode_count,
316                    injection_threats_detected: injection_count,
317                    total_scans: self.threats_blocked.load(Ordering::Relaxed),
318                }
319            },
320            Err(_) => crate::scanner::ScannerStats {
321                unicode_threats_detected: 0,
322                injection_threats_detected: 0,
323                total_scans: 0,
324            },
325        }
326    }
327
328    /// Set enabled state
329    pub fn set_enabled(&self, enabled: bool) {
330        // ShieldConfig is not mutable at runtime
331        // This would need to be handled differently
332        if enabled {
333            tracing::info!("Shield display enabled");
334            self.set_active(true);
335        } else {
336            tracing::info!("Shield display disabled");
337            self.set_active(false);
338        }
339    }
340}
kindly_guard_server/shield/mod.rs

kindly_guard_server/shield/
mod.rs
