Skip to content

Commit 8ad51fe

Browse files
committed
feat: allowlist support, presidio entity aliases, py.typed
Adds allowlist (exact values) and allowlist_patterns (full-match regexes) to scan/redact and threads them through both agent adapters: DATAFOG_HOOK_ALLOWLIST / DATAFOG_HOOK_ALLOWLIST_PATTERNS env vars for the Claude Code hook, allowlist/allowlist_patterns params for the LiteLLM guardrail. Motivated by a day of dogfooding: unix timestamps and numeric IDs match the PHONE pattern, and intentional identifiers (own support email, doc placeholders) should be exemptable. Accepts presidio-style entity names (EMAIL_ADDRESS, US_SSN) as input aliases via the existing canonical type map, ships a py.typed marker so downstream type checkers see our annotations, and backports the upstream-review fixes to the in-repo litellm adapter (guardrail spans recorded on the returned dict, redaction reported as intervention). Also corrects an entity-name documentation error introduced in #156: the scan API returns DATE and ZIP_CODE (DOB/ZIP are input aliases).
1 parent 5bf8805 commit 8ad51fe

10 files changed

Lines changed: 314 additions & 21 deletions

File tree

README.md

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,19 @@ values never echoed into logs or transcripts:
2929
```
3030

3131
Manual hook setup and limitations: [examples/claude_code_hook/](examples/claude_code_hook/).
32+
3233
- **LiteLLM guardrail** (`DataFogGuardrail`): redacts or blocks PII in
3334
requests and responses at the gateway, for any LiteLLM-proxied provider.
3435
In-process (~31µs per request), no sidecar service. Setup:
3536
[examples/litellm_guardrail/](examples/litellm_guardrail/).
3637

3738
Both default to the high-precision entity set (`EMAIL`, `PHONE`,
38-
`CREDIT_CARD`, `SSN`); noisier types are opt-in.
39+
`CREDIT_CARD`, `SSN`); noisier types are opt-in. Known-safe values can be
40+
exempted with an allowlist: `scan(text, allowlist=[...])` for exact values,
41+
`allowlist_patterns=[...]` for full-match regexes (e.g. `^\d{10}$` to stop
42+
unix timestamps matching as phone numbers) — available in both adapters and
43+
the API. Presidio-style entity names (`EMAIL_ADDRESS`, `PHONE_NUMBER`,
44+
`US_SSN`) are accepted as aliases for easy migration.
3945

4046
## Installation
4147

@@ -137,7 +143,7 @@ Use the engine that matches your accuracy and dependency constraints:
137143

138144
- `regex`:
139145
- Fastest and always available.
140-
- Best for default structured entities: `EMAIL`, `PHONE`, `SSN`, `CREDIT_CARD`, `IP_ADDRESS`, `DOB`, `ZIP`.
146+
- Best for default structured entities: `EMAIL`, `PHONE`, `SSN`, `CREDIT_CARD`, `IP_ADDRESS`, `DATE`, `ZIP_CODE` (`DOB` and `ZIP` are accepted as input aliases).
141147
- Use `locales=["de"]` for German structured IDs such as `DE_VAT_ID`, `DE_IBAN`, `DE_TAX_ID`, `DE_POSTAL_CODE`, and passport or residence permit numbers.
142148
- `spacy`:
143149
- Requires `pip install datafog[nlp]`.

datafog/__init__.py

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -153,14 +153,28 @@ def scan(
153153
engine: str = "regex",
154154
entity_types: list[str] | None = None,
155155
locales: list[str] | None = None,
156+
allowlist: list[str] | None = None,
157+
allowlist_patterns: list[str] | None = None,
156158
) -> ScanResult:
157159
"""
158160
v5-preview scan entrypoint.
159161
160162
Defaults to the lightweight regex engine so the core install works without
161163
optional dependency fallback warnings.
164+
165+
``allowlist`` exempts exact entity texts (your own support address, doc
166+
placeholders); ``allowlist_patterns`` exempts entities whose full text
167+
matches a regex (e.g. ``^\\d{10}$`` so unix timestamps stop matching as
168+
phone numbers).
162169
"""
163-
return _scan(text=text, engine=engine, entity_types=entity_types, locales=locales)
170+
return _scan(
171+
text=text,
172+
engine=engine,
173+
entity_types=entity_types,
174+
locales=locales,
175+
allowlist=allowlist,
176+
allowlist_patterns=allowlist_patterns,
177+
)
164178

165179

166180
def redact(
@@ -171,12 +185,17 @@ def redact(
171185
strategy: str = "token",
172186
preset: str | None = None,
173187
locales: list[str] | None = None,
188+
allowlist: list[str] | None = None,
189+
allowlist_patterns: list[str] | None = None,
174190
) -> RedactResult:
175191
"""
176192
v5-preview redaction entrypoint.
177193
178194
If entities are provided, redact those spans. Otherwise, scan text first
179-
using the selected engine and redact the detected entities.
195+
using the selected engine and redact the detected entities. ``allowlist``
196+
and ``allowlist_patterns`` exempt findings from redaction (exact text and
197+
full-text regex match respectively); they apply to the scan path and are
198+
rejected when explicit ``entities`` are supplied.
180199
"""
181200
if preset is not None:
182201
try:
@@ -186,6 +205,11 @@ def redact(
186205
raise ValueError(f"preset must be one of: {allowed}") from exc
187206

188207
if entities is not None:
208+
if allowlist or allowlist_patterns:
209+
raise ValueError(
210+
"allowlist/allowlist_patterns cannot be combined with explicit "
211+
"entities; filter the entities before calling redact"
212+
)
189213
return _redact_entities(text=text, entities=entities, strategy=strategy)
190214

191215
return _scan_and_redact(
@@ -194,6 +218,8 @@ def redact(
194218
entity_types=entity_types,
195219
strategy=strategy,
196220
locales=locales,
221+
allowlist=allowlist,
222+
allowlist_patterns=allowlist_patterns,
197223
)
198224

199225

datafog/engine.py

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from __future__ import annotations
44

55
import hashlib
6+
import re
67
import warnings
78
from dataclasses import dataclass
89
from functools import lru_cache
@@ -23,6 +24,9 @@
2324
"SOCIAL_SECURITY_NUMBER": "SSN",
2425
"CREDIT_CARD_NUMBER": "CREDIT_CARD",
2526
"DATE_OF_BIRTH": "DATE",
27+
# Presidio-compatible aliases, so configs migrate without renames.
28+
"EMAIL_ADDRESS": "EMAIL",
29+
"US_SSN": "SSN",
2630
}
2731

2832
ALL_ENTITY_TYPES = {
@@ -277,6 +281,42 @@ def _filter_entity_types(
277281
return [entity for entity in entities if entity.type in allowed]
278282

279283

284+
def _compile_allowlist_patterns(
285+
allowlist_patterns: Optional[list[str]],
286+
) -> list["re.Pattern[str]"]:
287+
compiled = []
288+
for raw in allowlist_patterns or []:
289+
try:
290+
compiled.append(re.compile(raw))
291+
except re.error as exc:
292+
raise ValueError(
293+
f"allowlist_patterns contains an invalid regex: {raw!r} ({exc})"
294+
) from None
295+
return compiled
296+
297+
298+
def _apply_allowlist(
299+
entities: list[Entity],
300+
allowlist: Optional[list[str]],
301+
allowlist_patterns: Optional[list[str]],
302+
) -> list[Entity]:
303+
"""Drop entities whose exact text is allowlisted.
304+
305+
Exact values match the full entity text; patterns must fullmatch it,
306+
so a partial match never suppresses a finding.
307+
"""
308+
if not allowlist and not allowlist_patterns:
309+
return entities
310+
exact = set(allowlist or [])
311+
patterns = _compile_allowlist_patterns(allowlist_patterns)
312+
return [
313+
entity
314+
for entity in entities
315+
if entity.text not in exact
316+
and not any(pattern.fullmatch(entity.text) for pattern in patterns)
317+
]
318+
319+
280320
def _needs_ner(entity_types: Optional[list[str]]) -> bool:
281321
if entity_types is None:
282322
return True
@@ -289,14 +329,25 @@ def scan(
289329
engine: str = "smart",
290330
entity_types: Optional[list[str]] = None,
291331
locales: Optional[list[str]] = None,
332+
allowlist: Optional[list[str]] = None,
333+
allowlist_patterns: Optional[list[str]] = None,
292334
) -> ScanResult:
293-
"""Scan text for PII entities."""
335+
"""Scan text for PII entities.
336+
337+
``allowlist`` exempts exact entity texts (e.g. your own support email);
338+
``allowlist_patterns`` exempts entities whose full text matches a regex
339+
(e.g. ``^\\d{10}$`` to stop unix timestamps matching as phone numbers).
340+
"""
294341
if not isinstance(text, str):
295342
raise TypeError("text must be a string")
296343

297344
if engine not in {"regex", "spacy", "gliner", "smart"}:
298345
raise ValueError("engine must be one of: regex, spacy, gliner, smart")
299346

347+
# Validate patterns up front so config errors fail fast even when the
348+
# text contains no entities.
349+
_compile_allowlist_patterns(allowlist_patterns)
350+
300351
regex_entities = _regex_entities(
301352
text,
302353
entity_types=entity_types,
@@ -305,6 +356,7 @@ def scan(
305356

306357
if engine == "regex":
307358
filtered = _filter_entity_types(regex_entities, entity_types)
359+
filtered = _apply_allowlist(filtered, allowlist, allowlist_patterns)
308360
return ScanResult(
309361
entities=_dedupe_entities(filtered), text=text, engine_used="regex"
310362
)
@@ -367,6 +419,7 @@ def scan(
367419
)
368420

369421
filtered = _filter_entity_types(combined, entity_types)
422+
filtered = _apply_allowlist(filtered, allowlist, allowlist_patterns)
370423
deduped = _dedupe_entities(filtered)
371424
return ScanResult(
372425
entities=deduped,
@@ -437,12 +490,16 @@ def scan_and_redact(
437490
entity_types: Optional[list[str]] = None,
438491
strategy: str = "token",
439492
locales: Optional[list[str]] = None,
493+
allowlist: Optional[list[str]] = None,
494+
allowlist_patterns: Optional[list[str]] = None,
440495
) -> RedactResult:
441496
"""Convenience wrapper: scan then redact."""
442497
scan_result = scan(
443498
text=text,
444499
engine=engine,
445500
entity_types=entity_types,
446501
locales=locales,
502+
allowlist=allowlist,
503+
allowlist_patterns=allowlist_patterns,
447504
)
448505
return redact(text=text, entities=scan_result.entities, strategy=strategy)

datafog/integrations/claude_code.py

Lines changed: 41 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@
1616
- ``DATAFOG_HOOK_ENTITIES``: comma-separated entity types to detect.
1717
Defaults to the high-precision set; noisy-in-code types (IP_ADDRESS,
1818
DOB, ZIP) must be opted into.
19+
- ``DATAFOG_HOOK_ALLOWLIST``: comma-separated exact values to exempt
20+
(your own support address, documentation placeholders).
21+
- ``DATAFOG_HOOK_ALLOWLIST_PATTERNS``: comma-separated regexes; findings
22+
whose full text matches are exempt (note: a pattern containing a comma
23+
cannot be expressed here).
1924
2025
Failure policy: fail open. A hook bug must never brick a Claude Code
2126
session, so any unexpected error exits non-blocking with no output.
@@ -59,6 +64,11 @@ def _action(env: Mapping[str, str]) -> str:
5964
return action if action in VALID_ACTIONS else "ask"
6065

6166

67+
def _csv_env(env: Mapping[str, str], name: str) -> list[str]:
68+
raw = env.get(name, "")
69+
return [item.strip() for item in raw.split(",") if item.strip()]
70+
71+
6272
def _iter_strings(value: Any) -> Iterator[str]:
6373
"""Yield every string embedded in a JSON-like structure.
6474
@@ -76,7 +86,12 @@ def _iter_strings(value: Any) -> Iterator[str]:
7686
stack.extend(current)
7787

7888

79-
def _scan_findings(value: Any, entity_types: list[str]) -> dict[str, int]:
89+
def _scan_findings(
90+
value: Any,
91+
entity_types: list[str],
92+
allowlist: list[str] | None = None,
93+
allowlist_patterns: list[str] | None = None,
94+
) -> dict[str, int]:
8095
"""Scan all strings in ``value``; return counts per entity type."""
8196
import datafog
8297

@@ -87,7 +102,13 @@ def _scan_findings(value: Any, entity_types: list[str]) -> dict[str, int]:
87102
break
88103
chunk = text[: min(MAX_SCAN_CHARS, total_budget)]
89104
total_budget -= len(chunk)
90-
result = datafog.scan(chunk, engine="regex", entity_types=entity_types)
105+
result = datafog.scan(
106+
chunk,
107+
engine="regex",
108+
entity_types=entity_types,
109+
allowlist=allowlist or None,
110+
allowlist_patterns=allowlist_patterns or None,
111+
)
91112
for entity in result.entities:
92113
counts[entity.type] = counts.get(entity.type, 0) + 1
93114
return counts
@@ -104,7 +125,12 @@ def _emit(event: str, fields: dict[str, Any]) -> str:
104125

105126

106127
def _handle_pre_tool_use(payload: dict, env: Mapping[str, str]) -> str:
107-
counts = _scan_findings(payload.get("tool_input"), _entity_types(env))
128+
counts = _scan_findings(
129+
payload.get("tool_input"),
130+
_entity_types(env),
131+
allowlist=_csv_env(env, "DATAFOG_HOOK_ALLOWLIST"),
132+
allowlist_patterns=_csv_env(env, "DATAFOG_HOOK_ALLOWLIST_PATTERNS"),
133+
)
108134
if not counts:
109135
return ""
110136
tool = payload.get("tool_name", "tool")
@@ -119,7 +145,12 @@ def _handle_pre_tool_use(payload: dict, env: Mapping[str, str]) -> str:
119145

120146

121147
def _handle_user_prompt_submit(payload: dict, env: Mapping[str, str]) -> str:
122-
counts = _scan_findings(payload.get("prompt"), _entity_types(env))
148+
counts = _scan_findings(
149+
payload.get("prompt"),
150+
_entity_types(env),
151+
allowlist=_csv_env(env, "DATAFOG_HOOK_ALLOWLIST"),
152+
allowlist_patterns=_csv_env(env, "DATAFOG_HOOK_ALLOWLIST_PATTERNS"),
153+
)
123154
if not counts:
124155
return ""
125156
context = (
@@ -130,7 +161,12 @@ def _handle_user_prompt_submit(payload: dict, env: Mapping[str, str]) -> str:
130161

131162

132163
def _handle_post_tool_use(payload: dict, env: Mapping[str, str]) -> str:
133-
counts = _scan_findings(payload.get("tool_response"), _entity_types(env))
164+
counts = _scan_findings(
165+
payload.get("tool_response"),
166+
_entity_types(env),
167+
allowlist=_csv_env(env, "DATAFOG_HOOK_ALLOWLIST"),
168+
allowlist_patterns=_csv_env(env, "DATAFOG_HOOK_ALLOWLIST_PATTERNS"),
169+
)
134170
if not counts:
135171
return ""
136172
tool = payload.get("tool_name", "tool")

0 commit comments

Comments
 (0)