assay_core/otel/
redaction.rs1use crate::config::otel::{PromptCaptureMode, RedactionConfig};
2use hmac::{Hmac, KeyInit, Mac};
3use sha2::Sha256;
4
5pub struct RedactionService {
6 mode: PromptCaptureMode,
7 config: RedactionConfig,
8 hmac_key: Vec<u8>, }
10
11impl RedactionService {
12 pub fn new(mode: PromptCaptureMode, config: RedactionConfig) -> Self {
13 let hmac_key = std::env::var("ASSAY_ORG_SECRET")
16 .unwrap_or_else(|_| "ephemeral-key".to_string())
17 .into_bytes();
18
19 Self {
20 mode,
21 config,
22 hmac_key,
23 }
24 }
25
26 pub fn should_capture(&self) -> bool {
28 !matches!(self.mode, PromptCaptureMode::Off)
29 }
30
31 pub fn is_blob_ref(&self) -> bool {
33 matches!(self.mode, PromptCaptureMode::BlobRef)
34 }
35
36 pub fn redact_inline(&self, content: &str) -> String {
39 let mut text = content.to_string();
40
41 text = self.scrub_control_chars(&text);
43
44 if text.trim_start().starts_with('{') {
46 if let Ok(mut v) = serde_json::from_str::<serde_json::Value>(&text) {
47 self.scrub_json(&mut v);
48 if let Ok(s) = serde_json::to_string(&v) {
49 text = s;
50 }
51 }
52 }
53
54 for policy in &self.config.policies {
57 if policy.starts_with("sk-") && text.contains(policy) {
60 text = text.replace(policy, "sk-[REDACTED]");
61 } else if text.contains("sk-") {
62 text = text.replace("sk-", "sk-[REDACTED]");
66 }
67 }
68
69 text
70 }
71
72 pub fn blob_ref(&self, content: &str) -> String {
75 type HmacSha256 = Hmac<Sha256>;
76 let mut mac =
77 HmacSha256::new_from_slice(&self.hmac_key).expect("HMAC can take key of any size");
78 mac.update(content.as_bytes());
79 let result = mac.finalize();
80 format!("hmac256:{}", hex::encode(result.into_bytes()))
82 }
83
84 fn scrub_control_chars(&self, input: &str) -> String {
85 input
87 .chars()
88 .filter(|c| {
89 let u = *c as u32;
90 u >= 32 || u == 10 || u == 13 || u == 9
91 })
92 .collect()
93 }
94
95 fn scrub_json(&self, v: &mut serde_json::Value) {
96 match v {
97 serde_json::Value::Object(map) => {
98 for (k, val) in map.iter_mut() {
99 if k == "api_key" || k == "authorization" || k == "token" {
100 *val = serde_json::Value::String("[REDACTED]".into());
101 } else {
102 self.scrub_json(val);
103 }
104 }
105 }
106 serde_json::Value::Array(arr) => {
107 for i in arr {
108 self.scrub_json(i);
109 }
110 }
111 _ => {}
112 }
113 }
114
115 pub fn pseudonymize(&self, id: &str) -> String {
117 self.blob_ref(id) }
119}