1use regex::Regex;
4use serde_json::Value;
5use std::collections::HashSet;
6
7pub const MAX_PREVIEW_LENGTH: usize = 100;
9
10const SENSITIVE_FIELDS: &[&str] = &[
12 "password",
13 "passwd",
14 "pwd",
15 "secret",
16 "token",
17 "api_key",
18 "apikey",
19 "access_token",
20 "refresh_token",
21 "auth_token",
22 "authorization",
23 "bearer",
24 "credit_card",
25 "card_number",
26 "cvv",
27 "ssn",
28 "social_security",
29 "private_key",
30 "privatekey",
31 "encryption_key",
32];
33
34lazy_static::lazy_static! {
35 static ref EMAIL_PATTERN: Regex = Regex::new(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b").unwrap();
37 static ref PHONE_PATTERN: Regex = Regex::new(r"\b(\+?\d{1,3}[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b").unwrap();
38 static ref CREDIT_CARD_PATTERN: Regex = Regex::new(r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b").unwrap();
39}
40
41pub fn truncate_string(text: &str, max_length: usize) -> String {
57 if text.len() <= max_length {
58 text.to_string()
59 } else {
60 format!("{}...", &text[..max_length])
61 }
62}
63
64pub fn sanitize_json(value: &Value) -> Value {
87 let sensitive_set: HashSet<&str> = SENSITIVE_FIELDS.iter().copied().collect();
88 sanitize_json_recursive(value, &sensitive_set)
89}
90
91fn sanitize_json_recursive(value: &Value, sensitive_fields: &HashSet<&str>) -> Value {
92 match value {
93 Value::Object(map) => {
94 let mut sanitized = serde_json::Map::new();
95 for (key, val) in map {
96 let key_lower = key.to_lowercase();
97 if sensitive_fields
98 .iter()
99 .any(|&field| key_lower.contains(field))
100 {
101 sanitized.insert(key.clone(), Value::String("[REDACTED]".to_string()));
102 } else {
103 sanitized.insert(key.clone(), sanitize_json_recursive(val, sensitive_fields));
104 }
105 }
106 Value::Object(sanitized)
107 }
108 Value::Array(arr) => Value::Array(
109 arr.iter()
110 .map(|v| sanitize_json_recursive(v, sensitive_fields))
111 .collect(),
112 ),
113 _ => value.clone(),
114 }
115}
116
117pub fn redact_pii(text: &str) -> String {
132 let mut result = text.to_string();
133
134 result = EMAIL_PATTERN.replace_all(&result, "[EMAIL]").to_string();
136
137 result = PHONE_PATTERN.replace_all(&result, "[PHONE]").to_string();
139
140 result = CREDIT_CARD_PATTERN
142 .replace_all(&result, "[CARD]")
143 .to_string();
144
145 result
146}
147
148pub fn safe_preview(text: &str, max_length: usize) -> String {
163 let redacted = redact_pii(text);
164 truncate_string(&redacted, max_length)
165}
166
167pub fn sanitize_tool_payload(payload: &Value, max_length: usize) -> String {
192 let sanitized_json = sanitize_json(payload);
193 let json_str = sanitized_json.to_string();
194 safe_preview(&json_str, max_length)
195}
196
197#[cfg(test)]
198mod tests {
199 use super::*;
200 use serde_json::json;
201
202 #[test]
203 fn test_truncate_string_short() {
204 let text = "Hello, world!";
205 assert_eq!(truncate_string(text, 100), "Hello, world!");
206 }
207
208 #[test]
209 fn test_truncate_string_long() {
210 let text = "a".repeat(150);
211 let truncated = truncate_string(&text, 100);
212 assert_eq!(truncated.len(), 103); assert!(truncated.ends_with("..."));
214 assert_eq!(&truncated[..100], &text[..100]);
215 }
216
217 #[test]
218 fn test_truncate_string_exact() {
219 let text = "a".repeat(100);
220 let truncated = truncate_string(&text, 100);
221 assert_eq!(truncated.len(), 100);
222 assert!(!truncated.ends_with("..."));
223 }
224
225 #[test]
226 fn test_sanitize_json_simple() {
227 let input = json!({
228 "username": "john",
229 "password": "secret123"
230 });
231
232 let sanitized = sanitize_json(&input);
233 assert_eq!(sanitized["username"], "john");
234 assert_eq!(sanitized["password"], "[REDACTED]");
235 }
236
237 #[test]
238 fn test_sanitize_json_nested() {
239 let input = json!({
240 "user": {
241 "name": "john",
242 "credentials": {
243 "password": "secret123",
244 "api_key": "sk-1234567890"
245 }
246 }
247 });
248
249 let sanitized = sanitize_json(&input);
250 assert_eq!(sanitized["user"]["name"], "john");
251 assert_eq!(sanitized["user"]["credentials"]["password"], "[REDACTED]");
252 assert_eq!(sanitized["user"]["credentials"]["api_key"], "[REDACTED]");
253 }
254
255 #[test]
256 fn test_sanitize_json_array() {
257 let input = json!({
258 "users": [
259 {"name": "john", "password": "secret1"},
260 {"name": "jane", "token": "abc123"}
261 ]
262 });
263
264 let sanitized = sanitize_json(&input);
265 assert_eq!(sanitized["users"][0]["name"], "john");
266 assert_eq!(sanitized["users"][0]["password"], "[REDACTED]");
267 assert_eq!(sanitized["users"][1]["name"], "jane");
268 assert_eq!(sanitized["users"][1]["token"], "[REDACTED]");
269 }
270
271 #[test]
272 fn test_sanitize_json_case_insensitive() {
273 let input = json!({
274 "Password": "secret123",
275 "API_KEY": "sk-1234567890",
276 "AccessToken": "token123"
277 });
278
279 let sanitized = sanitize_json(&input);
280 assert_eq!(sanitized["Password"], "[REDACTED]");
281 assert_eq!(sanitized["API_KEY"], "[REDACTED]");
282 assert_eq!(sanitized["AccessToken"], "[REDACTED]");
283 }
284
285 #[test]
286 fn test_redact_pii_email() {
287 let text = "Contact me at john.doe@example.com for more info";
288 let redacted = redact_pii(text);
289 assert!(redacted.contains("[EMAIL]"));
290 assert!(!redacted.contains("john.doe@example.com"));
291 }
292
293 #[test]
294 fn test_redact_pii_phone() {
295 let text = "Call me at 555-123-4567 or (555) 987-6543";
296 let redacted = redact_pii(text);
297 assert!(redacted.contains("[PHONE]"));
298 assert!(!redacted.contains("555-123-4567"));
299 assert!(!redacted.contains("555) 987-6543"));
300 }
301
302 #[test]
303 fn test_redact_pii_credit_card() {
304 let text = "Card number: 4532-1234-5678-9010";
305 let redacted = redact_pii(text);
306 assert!(redacted.contains("[CARD]"));
307 assert!(!redacted.contains("4532-1234-5678-9010"));
308 }
309
310 #[test]
311 fn test_redact_pii_multiple() {
312 let text = "Email: john@example.com, Phone: 555-123-1234, Card: 4532123456789010";
313 let redacted = redact_pii(text);
314 assert!(redacted.contains("[EMAIL]"));
315 assert!(redacted.contains("[PHONE]"));
316 assert!(redacted.contains("[CARD]"));
317 }
318
319 #[test]
320 fn test_safe_preview() {
321 let text = "My email is john@example.com and here's a very long message that goes on and on and on and on and on and on";
322 let preview = safe_preview(text, 50);
323
324 assert!(preview.len() <= 53); assert!(preview.contains("[EMAIL]"));
329 assert!(!preview.contains("john@example.com"));
330 }
331
332 #[test]
333 fn test_sanitize_tool_payload() {
334 let payload = json!({
335 "password": "secret123",
336 "api_key": "sk-1234567890",
337 "user": "john@example.com"
338 });
339
340 let sanitized = sanitize_tool_payload(&payload, 100);
341
342 assert!(
344 sanitized.len() <= 103,
345 "Length should be <= 103, got: {}",
346 sanitized.len()
347 );
348
349 assert!(
351 sanitized.contains("[REDACTED]"),
352 "Expected [REDACTED] in output, got: {}",
353 sanitized
354 );
355
356 assert!(
358 sanitized.contains("[EMAIL]"),
359 "Expected [EMAIL] in output, got: {}",
360 sanitized
361 );
362 }
363
364 #[test]
365 fn test_sanitize_tool_payload_long_message() {
366 let payload = json!({
367 "password": "secret123",
368 "message": "a".repeat(200)
369 });
370
371 let sanitized = sanitize_tool_payload(&payload, 100);
372
373 assert!(sanitized.len() <= 103);
375
376 assert!(sanitized.contains("[REDACTED]") || sanitized.ends_with("..."));
380 }
381
382 #[test]
383 fn test_sanitize_tool_payload_no_sensitive_data() {
384 let payload = json!({
385 "action": "get_weather",
386 "location": "Dubai"
387 });
388
389 let sanitized = sanitize_tool_payload(&payload, 100);
390 assert!(sanitized.contains("get_weather"));
391 assert!(sanitized.contains("Dubai"));
392 assert!(!sanitized.contains("[REDACTED]"));
393 }
394}