fleetforge_policy/
pii.rs

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/// Determines how the PII guardrail responds when sensitive data is detected.
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum PiiMode {
11    Redact,
12    Deny,
13    Allow,
14}
15
16/// Regex-based PII detector and redactor.
17#[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            // Raw string keeps the literal dot escaped for email domains.
29            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(&current, *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(&regexes, &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}