vtcode_core/exec/
pii_tokenizer.rs

1//! PII (Personally Identifiable Information) tokenization for data privacy.
2//!
3//! Automatically detects and tokenizes sensitive data before MCP tool calls,
4//! preventing PII from entering model context or being logged.
5//!
6//! Features:
7//! - Pattern-based detection (email, phone, SSN, credit card, etc.)
8//! - Secure token generation and storage
9//! - Automatic de-tokenization on tool result
10//! - Configurable patterns and policies
11//! - Audit trail of tokenized data
12
13use anyhow::{Context, Result};
14use hashbrown::HashMap;
15use once_cell::sync::Lazy;
16use regex::Regex;
17use serde::{Deserialize, Serialize};
18use std::sync::{Arc, Mutex};
19use tracing::debug;
20
21/// Types of PII that can be tokenized.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
23pub enum PiiType {
24    Email,
25    PhoneNumber,
26    SocialSecurityNumber,
27    CreditCard,
28    IpAddress,
29    ApiKey,
30    AuthToken,
31    Url,
32    Custom,
33}
34
35impl PiiType {
36    /// Return the PII type as a lowercase string identifier.
37    pub fn as_str(&self) -> &'static str {
38        match self {
39            Self::Email => "email",
40            Self::PhoneNumber => "phone_number",
41            Self::SocialSecurityNumber => "ssn",
42            Self::CreditCard => "credit_card",
43            Self::IpAddress => "ip_address",
44            Self::ApiKey => "api_key",
45            Self::AuthToken => "auth_token",
46            Self::Url => "url",
47            Self::Custom => "custom",
48        }
49    }
50}
51
52// Compile default PII patterns once to avoid repeated regex compilation overhead.
53static DEFAULT_PII_PATTERNS: Lazy<Result<Vec<(PiiType, Regex)>, String>> = Lazy::new(|| {
54    let patterns = vec![
55        (PiiType::Email, r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"),
56        (PiiType::PhoneNumber, r"(?:\+?1[-.\s]?)?\(?[0-9]{3}\)?[-.\s]?[0-9]{3}[-.\s]?[0-9]{4}"),
57        (PiiType::SocialSecurityNumber, r"[0-9]{3}-[0-9]{2}-[0-9]{4}"),
58        (PiiType::CreditCard, r"[0-9]{4}[\s-]?[0-9]{4}[\s-]?[0-9]{4}[\s-]?[0-9]{4}"),
59        (
60            PiiType::IpAddress,
61            r"(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)",
62        ),
63        (PiiType::ApiKey, r#"(?:api[_-]?key|apikey|API[_-]?KEY)\s*[:=]\s*['"]?[a-zA-Z0-9_-]{32,}['"]?"#),
64        (PiiType::AuthToken, r"(?:bearer|token|authorization)\s+[a-zA-Z0-9._-]+"),
65    ];
66
67    let mut compiled = Vec::with_capacity(patterns.len());
68    for (pii_type, pattern) in patterns {
69        match Regex::new(pattern) {
70            Ok(regex) => compiled.push((pii_type, regex)),
71            Err(e) => {
72                return Err(format!("Failed to compile PII regex for {pii_type:?}: {e}"));
73            }
74        }
75    }
76    Ok(compiled)
77});
78
79/// Detected PII instance with location and type.
80#[derive(Debug, Clone, Serialize, Deserialize)]
81pub struct DetectedPii {
82    pub value: String,
83    pub pii_type: PiiType,
84    pub start: usize,
85    pub end: usize,
86    pub context: String,
87}
88
89/// Token for replacing PII.
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct PiiToken {
92    pub token: String,
93    pub original_value: String,
94    pub pii_type: PiiType,
95    pub created_at: String,
96}
97
98/// Manager for PII tokenization with configurable detection patterns.
99#[derive(Clone)]
100pub struct PiiTokenizer {
101    patterns: HashMap<PiiType, Regex>,
102    inner: Arc<Mutex<PiiTokenizerInner>>,
103}
104
105/// Inner state for PiiTokenizer
106struct PiiTokenizerInner {
107    token_store: HashMap<String, PiiToken>,
108}
109
110impl PiiTokenizer {
111    /// Create a new PII tokenizer with default patterns.
112    pub fn new() -> Result<Self> {
113        // Build patterns from static defaults (compiled once)
114        // Note: Cloning Regex is cheap (Arc-based internally)
115        let patterns = DEFAULT_PII_PATTERNS
116            .as_ref()
117            .map_err(|e| anyhow::anyhow!("PII pattern initialization failed: {e}"))?
118            .iter()
119            .map(|(pii_type, regex)| (*pii_type, regex.clone()))
120            .collect();
121
122        Ok(Self {
123            patterns,
124            inner: Arc::new(Mutex::new(PiiTokenizerInner { token_store: HashMap::new() })),
125        })
126    }
127
128    /// Detect PII in a string.
129    pub fn detect_pii(&self, text: &str) -> Result<Vec<DetectedPii>> {
130        let mut detected = Vec::with_capacity(8);
131
132        for (pii_type, pattern) in &self.patterns {
133            for mat in pattern.find_iter(text) {
134                // Only allocate value when actually detected (lazy allocation)
135                let value = text[mat.start()..mat.end()].to_string();
136                let context_start = mat.start().saturating_sub(20);
137                let context_end = (mat.end() + 20).min(text.len());
138                let context = text[context_start..context_end].replace('\n', "\\n").replace('\r', "\\r");
139
140                debug!(
141                    pii_type = pii_type.as_str(),
142                    context = %context,
143                    "Detected PII in text"
144                );
145
146                detected.push(DetectedPii {
147                    value,
148                    pii_type: *pii_type,
149                    start: mat.start(),
150                    end: mat.end(),
151                    context,
152                });
153            }
154        }
155
156        Ok(detected)
157    }
158
159    /// Tokenize PII in a string, returning modified text and token map.
160    pub fn tokenize_string(&self, text: &str) -> Result<(String, HashMap<String, PiiToken>)> {
161        let mut detected = self.detect_pii(text)?;
162
163        if detected.is_empty() {
164            return Ok((text.to_string(), HashMap::new()));
165        }
166
167        // Ensure deterministic right-to-left replacement order by source offsets.
168        detected.sort_by_key(|d| d.start);
169
170        let mut result = text.to_string();
171        let mut new_tokens = HashMap::new();
172
173        // Process detections in reverse order to maintain offsets
174        for detection in detected.iter().rev() {
175            let token = self.generate_token(&detection.value, detection.pii_type)?;
176            let token_str = &token.token;
177            result.replace_range(detection.start..detection.end, token_str);
178            new_tokens.insert(token_str.clone(), token);
179        }
180
181        // Store tokens for later de-tokenization
182        {
183            let mut inner = self
184                .inner
185                .lock()
186                .map_err(|e| anyhow::anyhow!("Failed to acquire token store lock: {e}"))?;
187            // Clone tokens into the inner store so we can still return them
188            inner.token_store.extend(new_tokens.clone());
189        }
190
191        debug!(pii_count = detected.len(), "Tokenized PII in string");
192
193        Ok((result, new_tokens))
194    }
195
196    /// De-tokenize a string using stored token map.
197    pub fn detokenize_string(&self, text: &str) -> Result<String> {
198        let inner = self
199            .inner
200            .lock()
201            .map_err(|e| anyhow::anyhow!("Failed to acquire token store lock: {e}"))?;
202        let mut result = text.to_string();
203
204        for (token, pii_token) in inner.token_store.iter() {
205            result = result.replace(token, &pii_token.original_value);
206        }
207
208        Ok(result)
209    }
210
211    /// Clear all stored tokens (for security).
212    pub fn clear_tokens(&self) {
213        let mut inner = match self.inner.lock() {
214            Ok(guard) => guard,
215            Err(poisoned) => poisoned.into_inner(),
216        };
217        inner.token_store.clear();
218        debug!("Cleared all PII tokens");
219    }
220
221    /// Get audit trail of tokenized data.
222    pub fn audit_trail(&self) -> Result<Vec<(String, PiiType, String)>> {
223        let inner = self
224            .inner
225            .lock()
226            .map_err(|e| anyhow::anyhow!("Failed to acquire token store lock: {e}"))?;
227        Ok(inner
228            .token_store
229            .values()
230            .map(|t| (t.token.clone(), t.pii_type, t.created_at.clone()))
231            .collect())
232    }
233
234    /// Generate a secure token for PII value.
235    fn generate_token(&self, value: &str, pii_type: PiiType) -> Result<PiiToken> {
236        use std::collections::hash_map::DefaultHasher;
237        use std::hash::{Hash, Hasher};
238
239        let mut hasher = DefaultHasher::new();
240        value.hash(&mut hasher);
241        let hash = hasher.finish();
242
243        let token = format!("__PII_{}_{:x}__", pii_type.as_str(), hash);
244
245        Ok(PiiToken {
246            token,
247            original_value: value.to_string(),
248            pii_type,
249            created_at: chrono::Utc::now().to_rfc3339(),
250        })
251    }
252
253    /// Register custom PII pattern.
254    pub fn register_pattern(&mut self, pii_type: PiiType, pattern: &str) -> Result<()> {
255        let regex = Regex::new(pattern).context("invalid regex pattern for PII detection")?;
256        self.patterns.insert(pii_type, regex);
257        debug!(pii_type = pii_type.as_str(), pattern = pattern, "Registered custom PII pattern");
258        Ok(())
259    }
260}
261
262impl Default for PiiTokenizer {
263    fn default() -> Self {
264        Self::new().unwrap_or_else(|_| Self {
265            patterns: Default::default(),
266            inner: Arc::new(Mutex::new(PiiTokenizerInner { token_store: HashMap::new() })),
267        })
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274    use anyhow::Result;
275
276    #[test]
277    fn test_detect_email() -> Result<()> {
278        let tokenizer = PiiTokenizer::new()?;
279        let text = "Contact me at john@example.com for more info";
280        let detected = tokenizer.detect_pii(text)?;
281
282        assert!(!detected.is_empty());
283        assert!(detected.iter().any(|d| d.pii_type == PiiType::Email));
284        Ok(())
285    }
286
287    #[test]
288    fn test_detect_phone() -> Result<()> {
289        let tokenizer = PiiTokenizer::new()?;
290        let text = "Call me at 555-123-4567";
291        let detected = tokenizer.detect_pii(text)?;
292
293        assert!(!detected.is_empty());
294        assert!(detected.iter().any(|d| d.pii_type == PiiType::PhoneNumber));
295        Ok(())
296    }
297
298    #[test]
299    fn test_tokenize_string() -> Result<()> {
300        let tokenizer = PiiTokenizer::new()?;
301        let text = "Email: john@example.com, Phone: 555-123-4567";
302        let (tokenized, tokens) = tokenizer.tokenize_string(text)?;
303
304        assert!(tokenized.contains("__PII_"));
305        assert!(!tokenized.contains("john@example.com"));
306        assert!(!tokens.is_empty());
307        Ok(())
308    }
309
310    #[test]
311    fn test_no_pii_detected() -> Result<()> {
312        let tokenizer = PiiTokenizer::new()?;
313        let text = "This is regular text with no sensitive information";
314        let detected = tokenizer.detect_pii(text)?;
315
316        assert!(detected.is_empty());
317        Ok(())
318    }
319}
vtcode_core/exec/pii_tokenizer.rs

vtcode_core/exec/
pii_tokenizer.rs
