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.get(pattern_type).copied().unwrap_or(0)
65    }
66
67    pub fn get_most_detected_pattern(&self) -> Option<&str> {
68        self.pattern_distribution
69            .iter()
70            .max_by_key(|(_, count)| *count)
71            .map(|(name, _)| name.as_str())
72    }
73
74    pub fn get_high_severity_events(&self) -> Vec<AuditEvent> {
75        self.audit_events
76            .iter()
77            .filter(|e| e.severity == "high" || e.severity == "critical")
78            .cloned()
79            .collect()
80    }
81}
82
83impl Default for SecurityMetrics {
84    fn default() -> Self {
85        Self::new()
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[test]
94    fn test_record_detection() {
95        let mut metrics = SecurityMetrics::new();
96        metrics.record_detection("email".to_owned());
97        metrics.record_detection("ssn".to_owned());
98        metrics.record_detection("email".to_owned());
99
100        assert_eq!(metrics.pii_detections, 3);
101        assert_eq!(metrics.get_pattern_count("email"), 2);
102        assert_eq!(metrics.get_pattern_count("ssn"), 1);
103    }
104
105    #[test]
106    fn test_record_tokenization() {
107        let mut metrics = SecurityMetrics::new();
108        metrics.record_tokenization(5);
109        metrics.record_tokenization(3);
110
111        assert_eq!(metrics.tokens_created, 8);
112    }
113
114    #[test]
115    fn test_record_audit_event() {
116        let mut metrics = SecurityMetrics::new();
117        metrics.record_audit_event("pii_detected".to_owned(), "info".to_owned());
118        metrics.record_audit_event("tokenized".to_owned(), "info".to_owned());
119
120        assert_eq!(metrics.audit_events.len(), 2);
121    }
122
123    #[test]
124    fn test_detection_rate() {
125        let mut metrics = SecurityMetrics::new();
126        metrics.record_detection("email".to_owned());
127        metrics.record_detection("email".to_owned());
128        metrics.record_audit_event("test".to_owned(), "info".to_owned());
129
130        let rate = metrics.detection_rate();
131        assert!(rate > 0.0);
132    }
133
134    #[test]
135    fn test_most_detected_pattern() {
136        let mut metrics = SecurityMetrics::new();
137        metrics.record_detection("email".to_owned());
138        metrics.record_detection("email".to_owned());
139        metrics.record_detection("email".to_owned());
140        metrics.record_detection("ssn".to_owned());
141
142        assert_eq!(metrics.get_most_detected_pattern(), Some("email"));
143    }
144
145    #[test]
146    fn test_high_severity_events() {
147        let mut metrics = SecurityMetrics::new();
148        metrics.record_audit_event("event1".to_owned(), "low".to_owned());
149        metrics.record_audit_event("event2".to_owned(), "high".to_owned());
150        metrics.record_audit_event("event3".to_owned(), "critical".to_owned());
151        metrics.record_audit_event("event4".to_owned(), "info".to_owned());
152
153        let high_severity = metrics.get_high_severity_events();
154        assert_eq!(high_severity.len(), 2);
155    }
156
157    #[test]
158    fn test_audit_events_limit() {
159        let mut metrics = SecurityMetrics::new();
160        for i in 0..150 {
161            metrics.record_audit_event(format!("event_{i}"), "info".to_owned());
162        }
163
164        assert_eq!(metrics.audit_events.len(), 100);
165    }
166}
vtcode_core/metrics/security_metrics.rs

vtcode_core/metrics/
security_metrics.rs
