1use anyhow::Result;
2use async_trait::async_trait;
3use regex::Regex;
4use serde_json::{json, Value};
5
6use crate::{Decision, DecisionEffect, PolicyEngine, PolicyRequest};
7
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum PiiMode {
11 Redact,
12 Deny,
13 Allow,
14}
15
16#[derive(Clone)]
18pub struct BasicPiiPolicy {
19 email: Regex,
20 phone: Regex,
21 card: Regex,
22 mode: PiiMode,
23}
24
25impl BasicPiiPolicy {
26 pub fn new(mode: PiiMode) -> Self {
27 Self {
28 email: Regex::new(r"(?i)[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}")
30 .expect("valid email regex"),
31 phone: Regex::new(
32 r"(?x)\b(?:\+?\d{1,3}[-. ]?)?(?:\(?\d{3}\)?[-. ]?)?\d{3}[-. ]?\d{4}\b",
33 )
34 .expect("valid phone regex"),
35 card: Regex::new(r"\b(?:\d[ -]?){13,16}\b").expect("valid card regex"),
36 mode,
37 }
38 }
39
40 pub fn with_mode(mut self, mode: PiiMode) -> Self {
41 self.mode = mode;
42 self
43 }
44}
45
46impl Default for BasicPiiPolicy {
47 fn default() -> Self {
48 Self::new(PiiMode::Redact)
49 }
50}
51
52fn redact_value(regexes: &[(&Regex, &str)], value: &Value, redacted: &mut bool) -> Value {
53 match value {
54 Value::String(s) => {
55 let mut current = s.clone();
56 for (regex, replacement) in regexes {
57 let next = regex.replace_all(¤t, *replacement).to_string();
58 if next != current {
59 *redacted = true;
60 current = next;
61 }
62 }
63 Value::String(current)
64 }
65 Value::Array(items) => {
66 let mut changed = false;
67 let redacted_items = items
68 .iter()
69 .map(|item| redact_value(regexes, item, &mut changed))
70 .collect();
71 if changed {
72 *redacted = true;
73 }
74 Value::Array(redacted_items)
75 }
76 Value::Object(map) => {
77 let mut changed = false;
78 let mut new_map = serde_json::Map::with_capacity(map.len());
79 for (k, v) in map {
80 let next = redact_value(regexes, v, &mut changed);
81 new_map.insert(k.clone(), next);
82 }
83 if changed {
84 *redacted = true;
85 }
86 Value::Object(new_map)
87 }
88 _ => value.clone(),
89 }
90}
91
92fn guardrails(request: &PolicyRequest) -> Vec<String> {
93 request
94 .context()
95 .get("policy")
96 .and_then(Value::as_object)
97 .and_then(|obj| obj.get("guardrails"))
98 .and_then(Value::as_array)
99 .map(|arr| {
100 arr.iter()
101 .filter_map(Value::as_str)
102 .map(|s| s.to_string())
103 .collect()
104 })
105 .unwrap_or_default()
106}
107
108fn payload_inputs(request: &PolicyRequest, payload: &Value) -> Value {
109 payload
110 .get("inputs")
111 .cloned()
112 .or_else(|| request.context().get("inputs").cloned())
113 .unwrap_or(Value::Null)
114}
115
116fn payload_output(payload: &Value) -> Value {
117 payload.get("output").cloned().unwrap_or(Value::Null)
118}
119
120#[async_trait]
121impl PolicyEngine for BasicPiiPolicy {
122 async fn evaluate(&self, request: &PolicyRequest) -> Result<Decision> {
123 let regexes = [
124 (&self.email, "[redacted-email]"),
125 (&self.phone, "[redacted-phone]"),
126 (&self.card, "[redacted-card]"),
127 ];
128
129 let payload = request
130 .context()
131 .get("payload")
132 .cloned()
133 .unwrap_or(Value::Null);
134 let boundary = request
135 .context()
136 .get("boundary")
137 .and_then(Value::as_str)
138 .unwrap_or_default()
139 .to_string();
140
141 let (target, patch_path) = match boundary.as_str() {
142 "egress_prompt" | "egress_tool" | "egress_memory" => {
143 (payload_output(&payload), "/output")
144 }
145 _ => (payload_inputs(request, &payload), "/inputs"),
146 };
147
148 let mut redacted_flag = false;
149 let redacted_value = redact_value(®exes, &target, &mut redacted_flag);
150
151 if !redacted_flag {
152 return Ok(Decision {
153 effect: DecisionEffect::Allow,
154 reason: None,
155 patches: Vec::new(),
156 });
157 }
158
159 let guardrails = guardrails(request);
160 let effective_mode = if guardrails.iter().any(|rule| rule == "deny_on_pii") {
161 PiiMode::Deny
162 } else if guardrails.iter().any(|rule| rule == "skip_redaction") {
163 PiiMode::Allow
164 } else {
165 self.mode
166 };
167
168 match effective_mode {
169 PiiMode::Allow => Ok(Decision {
170 effect: DecisionEffect::Allow,
171 reason: Some("PII detected but redaction skipped".into()),
172 patches: Vec::new(),
173 }),
174 PiiMode::Deny => Ok(Decision {
175 effect: DecisionEffect::Deny,
176 reason: Some("PII detected".into()),
177 patches: Vec::new(),
178 }),
179 PiiMode::Redact => {
180 let patch = json!({
181 "op": "replace",
182 "path": patch_path,
183 "value": redacted_value,
184 });
185 Ok(Decision {
186 effect: DecisionEffect::Redact,
187 reason: Some("PII redacted".into()),
188 patches: vec![patch],
189 })
190 }
191 }
192 }
193}