vtcode_core/metrics/
security_metrics.rs

1use chrono::{DateTime, Utc};
2use hashbrown::HashMap;
3use serde::{Deserialize, Serialize};
4use std::collections::VecDeque;
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct SecurityMetrics {
8    pub pii_detections: u64,
9    pub tokens_created: u64,
10    pub pattern_distribution: HashMap<String, u64>,
11    pub audit_events: VecDeque<AuditEvent>,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct AuditEvent {
16    pub event_type: String,
17    pub pattern_type: String,
18    pub severity: String,
19    pub timestamp: DateTime<Utc>,
20}
21
22impl SecurityMetrics {
23    pub fn new() -> Self {
24        Self {
25            pii_detections: 0,
26            tokens_created: 0,
27            pattern_distribution: HashMap::new(),
28            audit_events: VecDeque::with_capacity(100),
29        }
30    }
31
32    pub fn record_detection(&mut self, pattern_type: String) {
33        self.pii_detections += 1;
34        *self.pattern_distribution.entry(pattern_type).or_insert(0) += 1;
35    }
36
37    pub fn record_tokenization(&mut self, token_count: usize) {
38        self.tokens_created += token_count as u64;
39    }
40
41    pub fn record_audit_event(&mut self, event_type: String, severity: String) {
42        let event = AuditEvent {
43            event_type,
44            pattern_type: String::new(),
45            severity,
46            timestamp: Utc::now(),
47        };
48
49        if self.audit_events.len() >= 100 {
50            self.audit_events.pop_front();
51        }
52        self.audit_events.push_back(event);
53    }
54
55    pub fn detection_rate(&self) -> f64 {
56        if !self.audit_events.is_empty() {
57            self.pii_detections as f64 / self.audit_events.len() as f64
58        } else {
59            0.0
60        }
61    }
62
63    pub fn get_pattern_count(&self, pattern_type: &str) -> u64 {
64        self.pattern_distribution
65            .get(pattern_type)
66            .copied()
67            .unwrap_or(0)
68    }
69
70    pub fn get_most_detected_pattern(&self) -> Option<&str> {
71        self.pattern_distribution
72            .iter()
73            .max_by_key(|(_, count)| *count)
74            .map(|(name, _)| name.as_str())
75    }
76
77    pub fn get_high_severity_events(&self) -> Vec<AuditEvent> {
78        self.audit_events
79            .iter()
80            .filter(|e| e.severity == "high" || e.severity == "critical")
81            .cloned()
82            .collect()
83    }
84}
85
86impl Default for SecurityMetrics {
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_record_detection() {
98        let mut metrics = SecurityMetrics::new();
99        metrics.record_detection("email".to_owned());
100        metrics.record_detection("ssn".to_owned());
101        metrics.record_detection("email".to_owned());
102
103        assert_eq!(metrics.pii_detections, 3);
104        assert_eq!(metrics.get_pattern_count("email"), 2);
105        assert_eq!(metrics.get_pattern_count("ssn"), 1);
106    }
107
108    #[test]
109    fn test_record_tokenization() {
110        let mut metrics = SecurityMetrics::new();
111        metrics.record_tokenization(5);
112        metrics.record_tokenization(3);
113
114        assert_eq!(metrics.tokens_created, 8);
115    }
116
117    #[test]
118    fn test_record_audit_event() {
119        let mut metrics = SecurityMetrics::new();
120        metrics.record_audit_event("pii_detected".to_owned(), "info".to_owned());
121        metrics.record_audit_event("tokenized".to_owned(), "info".to_owned());
122
123        assert_eq!(metrics.audit_events.len(), 2);
124    }
125
126    #[test]
127    fn test_detection_rate() {
128        let mut metrics = SecurityMetrics::new();
129        metrics.record_detection("email".to_owned());
130        metrics.record_detection("email".to_owned());
131        metrics.record_audit_event("test".to_owned(), "info".to_owned());
132
133        let rate = metrics.detection_rate();
134        assert!(rate > 0.0);
135    }
136
137    #[test]
138    fn test_most_detected_pattern() {
139        let mut metrics = SecurityMetrics::new();
140        metrics.record_detection("email".to_owned());
141        metrics.record_detection("email".to_owned());
142        metrics.record_detection("email".to_owned());
143        metrics.record_detection("ssn".to_owned());
144
145        assert_eq!(metrics.get_most_detected_pattern(), Some("email"));
146    }
147
148    #[test]
149    fn test_high_severity_events() {
150        let mut metrics = SecurityMetrics::new();
151        metrics.record_audit_event("event1".to_owned(), "low".to_owned());
152        metrics.record_audit_event("event2".to_owned(), "high".to_owned());
153        metrics.record_audit_event("event3".to_owned(), "critical".to_owned());
154        metrics.record_audit_event("event4".to_owned(), "info".to_owned());
155
156        let high_severity = metrics.get_high_severity_events();
157        assert_eq!(high_severity.len(), 2);
158    }
159
160    #[test]
161    fn test_audit_events_limit() {
162        let mut metrics = SecurityMetrics::new();
163        for i in 0..150 {
164            metrics.record_audit_event(format!("event_{}", i), "info".to_owned());
165        }
166
167        assert_eq!(metrics.audit_events.len(), 100);
168    }
169}
vtcode_core/metrics/security_metrics.rs

vtcode_core/metrics/
security_metrics.rs
