feat: v0.14.0 — live PII and sensitive data masking in proxy layer (#30)

* feat: live PII and sensitive data masking in proxy layer

Adds masking.py with MaskingConfig and mask_event_data() that applies
PII detection (emails, phones, credit cards, SSNs, AWS ARNs, custom
patterns) on top of the existing secret redaction. Both the stdio proxy
and HTTP proxy now accept a masking_config parameter and run masking
before storing events.

Closes #20

Co-authored-by: Ona <no-reply@ona.com>

* fix: reduce phone regex false positives and fix ARN account ID masking

- Phone regex now requires a separator (dash, dot, or parens) around
  the area code, preventing false positives on IP addresses and plain
  numeric sequences like '192.168.1.1' or version strings
- AWS ARN masking now uses str.replace(..., 1) to avoid replacing the
  account ID string if it appears elsewhere in the ARN path

Co-authored-by: Ona <no-reply@ona.com>

---------

Co-authored-by: Ona <ona@gitpod.io>
Co-authored-by: Ona <no-reply@ona.com>

Author: 3 people

src/agent_trace/http_proxy.py

@@ -28,7 +28,7 @@
 
 from .models import EventType, SessionMeta, TraceEvent
 from .proxy import _classify_message
-from .redact import redact_data
+from .masking import MaskingConfig, mask_event_data
 from .store import TraceStore
 
 
@@ -41,13 +41,18 @@ class _ProxyHandler(BaseHTTPRequestHandler):
     meta: SessionMeta | None = None
     on_event: Callable[[TraceEvent], None] | None = None
     redact: bool = False
+    masking_config: MaskingConfig | None = None
     pending_calls: dict[Any, TraceEvent] = {}
 
     def _emit(self, event: TraceEvent) -> None:
         if self.meta:
             event.session_id = self.meta.session_id
-        if self.redact:
-            event.data = redact_data(event.data)
+        if self.redact or self.masking_config:
+            event.data = mask_event_data(
+                event.data,
+                config=self.masking_config,
+                redact_secrets=self.redact,
+            )
         if self.store and self.meta:
             self.store.append_event(self.meta.session_id, event)
 

src/agent_trace/masking.py

@@ -0,0 +1,205 @@
+"""Live PII and sensitive data masking for the proxy layer.
+
+Extends redact.py with PII detection (emails, phone numbers, credit cards,
+SSNs, IP addresses, names in common patterns) applied in real-time as events
+flow through the proxy. Operates on raw JSON-RPC message bodies before they
+are stored as trace events.
+
+Usage in proxy:
+    from .masking import mask_event_data
+    event.data = mask_event_data(event.data, config)
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+from .redact import redact_data, REDACTED
+
+# ---------------------------------------------------------------------------
+# PII patterns
+# ---------------------------------------------------------------------------
+
+# Email addresses
+_EMAIL_RE = re.compile(
+    r"\b[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}\b"
+)
+
+# Phone numbers (US and international formats)
+# Require a separator (space, dash, dot, or parens) to reduce false positives
+# on plain numeric sequences like IP addresses or version numbers.
+_PHONE_RE = re.compile(
+    r"(?<!\d)(?:\+?1[\s\-.]?)?"
+    r"(?:\(\d{3}\)[\s\-.]|\d{3}[\-.])"  # area code must have separator
+    r"\d{3}[\s\-.]?\d{4}"
+    r"(?!\d)"
+)
+
+# Credit card numbers (Visa, MC, Amex, Discover — 13-16 digits with optional separators)
+_CC_RE = re.compile(
+    r"\b(?:4[0-9]{12}(?:[0-9]{3})?"           # Visa
+    r"|5[1-5][0-9]{14}"                         # Mastercard
+    r"|3[47][0-9]{13}"                          # Amex
+    r"|6(?:011|5[0-9]{2})[0-9]{12}"            # Discover
+    r"|(?:\d{4}[\s\-]){3}\d{4})\b"             # Generic 16-digit with separators
+)
+
+# US Social Security Numbers
+_SSN_RE = re.compile(
+    r"\b(?!000|666|9\d{2})\d{3}[\s\-]"
+    r"(?!00)\d{2}[\s\-]"
+    r"(?!0000)\d{4}\b"
+)
+
+# IPv4 addresses (private ranges are not masked; public IPs are)
+_IPV4_RE = re.compile(
+    r"\b(?!(?:10|127|172\.(?:1[6-9]|2\d|3[01])|192\.168)\.\d+\.\d+)"
+    r"(?:\d{1,3}\.){3}\d{1,3}\b"
+)
+
+# AWS account IDs (12-digit numbers in ARNs)
+_AWS_ARN_RE = re.compile(r"arn:aws:[a-z0-9\-]+:[a-z0-9\-]*:(\d{12}):")
+
+# Generic UUIDs that look like user/account IDs (not session IDs)
+# Only mask when they appear as values of keys containing "user", "account", "customer"
+_UUID_RE = re.compile(
+    r"\b[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\b",
+    re.IGNORECASE,
+)
+
+_USER_KEYS = {
+    "user_id", "userid", "user-id", "account_id", "accountid",
+    "customer_id", "customerid", "person_id", "personid",
+    "email", "phone", "mobile", "ssn", "dob", "date_of_birth",
+    "credit_card", "card_number", "cvv", "billing_address",
+}
+
+# ---------------------------------------------------------------------------
+# Masking config
+# ---------------------------------------------------------------------------
+
+@dataclass
+class MaskingConfig:
+    mask_emails: bool = True
+    mask_phones: bool = True
+    mask_credit_cards: bool = True
+    mask_ssn: bool = True
+    mask_public_ips: bool = False   # off by default — too noisy
+    mask_aws_arns: bool = True
+    mask_user_id_keys: bool = True
+    # Custom regex patterns provided by the user
+    custom_patterns: list[str] = field(default_factory=list)
+    # Replacement token (default: [MASKED])
+    replacement: str = "[MASKED]"
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "MaskingConfig":
+        return cls(
+            mask_emails=d.get("emails", True),
+            mask_phones=d.get("phones", True),
+            mask_credit_cards=d.get("credit_cards", True),
+            mask_ssn=d.get("ssn", True),
+            mask_public_ips=d.get("public_ips", False),
+            mask_aws_arns=d.get("aws_arns", True),
+            mask_user_id_keys=d.get("user_id_keys", True),
+            custom_patterns=d.get("custom_patterns", []),
+            replacement=d.get("replacement", "[MASKED]"),
+        )
+
+    @classmethod
+    def default(cls) -> "MaskingConfig":
+        return cls()
+
+
+# ---------------------------------------------------------------------------
+# String-level masking
+# ---------------------------------------------------------------------------
+
+def _mask_string(value: str, config: MaskingConfig) -> str:
+    """Apply all enabled PII patterns to a string value."""
+    result = value
+    repl = config.replacement
+
+    if config.mask_emails:
+        result = _EMAIL_RE.sub(repl, result)
+
+    if config.mask_phones:
+        result = _PHONE_RE.sub(repl, result)
+
+    if config.mask_credit_cards:
+        result = _CC_RE.sub(repl, result)
+
+    if config.mask_ssn:
+        result = _SSN_RE.sub(repl, result)
+
+    if config.mask_public_ips:
+        result = _IPV4_RE.sub(repl, result)
+
+    if config.mask_aws_arns:
+        # Replace only the account ID portion (group 1) within the ARN
+        def _mask_arn(m: re.Match) -> str:
+            return m.group(0).replace(m.group(1), repl, 1)
+        result = _AWS_ARN_RE.sub(_mask_arn, result)
+
+    for pat_str in config.custom_patterns:
+        try:
+            result = re.sub(pat_str, repl, result)
+        except re.error:
+            pass
+
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Recursive data masking
+# ---------------------------------------------------------------------------
+
+def mask_data(data: Any, config: MaskingConfig, parent_key: str = "") -> Any:
+    """Recursively apply PII masking to event data.
+
+    Runs *after* redact_data (which handles secrets/tokens). This layer
+    handles PII: emails, phones, credit cards, SSNs, etc.
+    """
+    if isinstance(data, dict):
+        result = {}
+        for k, v in data.items():
+            k_lower = k.lower().strip()
+            if config.mask_user_id_keys and k_lower in _USER_KEYS:
+                result[k] = config.replacement if isinstance(v, str) else v
+            else:
+                result[k] = mask_data(v, config, parent_key=k)
+        return result
+
+    if isinstance(data, list):
+        return [mask_data(item, config, parent_key=parent_key) for item in data]
+
+    if isinstance(data, str):
+        return _mask_string(data, config)
+
+    return data
+
+
+# ---------------------------------------------------------------------------
+# Combined masking entry point (secrets + PII)
+# ---------------------------------------------------------------------------
+
+def mask_event_data(
+    data: Any,
+    config: MaskingConfig | None = None,
+    redact_secrets: bool = True,
+) -> Any:
+    """Apply secret redaction and PII masking to event data.
+
+    This is the single entry point used by the proxy layer. It runs
+    redact_data first (for API keys, tokens, etc.) then mask_data (for PII).
+    """
+    if config is None:
+        config = MaskingConfig.default()
+
+    result = data
+    if redact_secrets:
+        result = redact_data(result)
+    result = mask_data(result, config)
+    return result

src/agent_trace/proxy.py

@@ -28,7 +28,7 @@
 from typing import IO, Any, Callable
 
 from .models import EventType, SessionMeta, TraceEvent
-from .redact import redact_data
+from .masking import MaskingConfig, mask_event_data
 from .store import TraceStore
 
 
@@ -202,18 +202,24 @@ def __init__(
         session_meta: SessionMeta,
         on_event: Callable[[TraceEvent], None] | None = None,
         redact: bool = False,
+        masking_config: MaskingConfig | None = None,
     ):
         self.server_command = server_command
         self.store = store
         self.meta = session_meta
         self.on_event = on_event
         self.redact = redact
+        self.masking_config = masking_config
         self._pending_calls: dict[Any, TraceEvent] = {}
 
     def _emit(self, event: TraceEvent) -> None:
         event.session_id = self.meta.session_id
-        if self.redact:
-            event.data = redact_data(event.data)
+        if self.redact or self.masking_config:
+            event.data = mask_event_data(
+                event.data,
+                config=self.masking_config,
+                redact_secrets=self.redact,
+            )
         self.store.append_event(self.meta.session_id, event)
 
         # update counters

tests/test_masking.py

@@ -0,0 +1,110 @@
+"""Tests for live PII masking (issue #20)."""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+
+from agent_trace.masking import MaskingConfig, mask_data, mask_event_data, _mask_string
+
+
+class TestMaskString(unittest.TestCase):
+    def setUp(self):
+        self.cfg = MaskingConfig.default()
+
+    def test_masks_email(self):
+        result = _mask_string("Contact alice@example.com for help", self.cfg)
+        self.assertNotIn("alice@example.com", result)
+        self.assertIn("[MASKED]", result)
+
+    def test_masks_phone_us(self):
+        result = _mask_string("Call 555-867-5309 now", self.cfg)
+        self.assertNotIn("555-867-5309", result)
+
+    def test_masks_credit_card(self):
+        result = _mask_string("Card: 4111 1111 1111 1111", self.cfg)
+        self.assertNotIn("4111 1111 1111 1111", result)
+
+    def test_masks_ssn(self):
+        result = _mask_string("SSN: 123-45-6789", self.cfg)
+        self.assertNotIn("123-45-6789", result)
+
+    def test_safe_string_unchanged(self):
+        result = _mask_string("hello world", self.cfg)
+        self.assertEqual(result, "hello world")
+
+    def test_custom_pattern(self):
+        cfg = MaskingConfig(custom_patterns=[r"\bACCT-\d+\b"])
+        result = _mask_string("Account ACCT-12345 is active", cfg)
+        self.assertNotIn("ACCT-12345", result)
+
+    def test_public_ip_off_by_default(self):
+        cfg = MaskingConfig(mask_public_ips=False)
+        result = _mask_string("Server at 8.8.8.8", cfg)
+        self.assertIn("8.8.8.8", result)
+
+    def test_public_ip_masked_when_enabled(self):
+        cfg = MaskingConfig(mask_public_ips=True)
+        result = _mask_string("Server at 8.8.8.8", cfg)
+        self.assertNotIn("8.8.8.8", result)
+
+
+class TestMaskData(unittest.TestCase):
+    def setUp(self):
+        self.cfg = MaskingConfig.default()
+
+    def test_masks_email_in_dict_value(self):
+        data = {"message": "Send to bob@test.org please"}
+        result = mask_data(data, self.cfg)
+        self.assertNotIn("bob@test.org", result["message"])
+
+    def test_masks_sensitive_key(self):
+        data = {"email": "alice@example.com", "name": "Alice"}
+        result = mask_data(data, self.cfg)
+        self.assertEqual(result["email"], "[MASKED]")
+        self.assertEqual(result["name"], "Alice")
+
+    def test_masks_nested(self):
+        data = {"user": {"contact": "call 555-123-4567"}}
+        result = mask_data(data, self.cfg)
+        self.assertNotIn("555-123-4567", result["user"]["contact"])
+
+    def test_masks_list_items(self):
+        data = {"notes": ["email: test@example.com", "safe note"]}
+        result = mask_data(data, self.cfg)
+        self.assertNotIn("test@example.com", result["notes"][0])
+        self.assertEqual(result["notes"][1], "safe note")
+
+    def test_non_string_unchanged(self):
+        data = {"count": 42, "flag": True}
+        result = mask_data(data, self.cfg)
+        self.assertEqual(result["count"], 42)
+        self.assertEqual(result["flag"], True)
+
+
+class TestMaskEventData(unittest.TestCase):
+    def test_combined_secrets_and_pii(self):
+        data = {
+            "api_key": "sk-abc123def456ghi789jkl012mno345pqr678",
+            "message": "Contact alice@example.com",
+        }
+        result = mask_event_data(data)
+        # Secret redacted
+        self.assertNotIn("sk-abc", str(result.get("api_key", "")))
+        # PII masked
+        self.assertNotIn("alice@example.com", result.get("message", ""))
+
+    def test_no_redact_secrets_flag(self):
+        data = {"message": "email: test@example.com"}
+        result = mask_event_data(data, redact_secrets=False)
+        self.assertNotIn("test@example.com", result["message"])
+
+    def test_none_config_uses_defaults(self):
+        data = {"note": "phone: 555-123-4567"}
+        result = mask_event_data(data, config=None)
+        self.assertNotIn("555-123-4567", result["note"])
+
+
+if __name__ == "__main__":
+    unittest.main()