Masking PII in DR Sandboxes with Python

This page implements the transform-and-scan gate defined in drill data masking and anonymization: a headless Python tool that connects to a freshly restored PostgreSQL sandbox, rewrites every configured PII column with a deterministic HMAC transform, batches the UPDATEs so a hundred-million-row table does not blow out the drill’s time budget, then runs an independent residual-PII scan and returns a POSIX exit code the orchestrator branches on. It runs only against the isolated instance that sandbox provisioning stands up, assumes the network isolation from the security boundaries for DR environments reference is already in force, and its clean-scan verdict is the precondition that lets smoke-test routing logic publish the sandbox for validation. Determinism is the load-bearing property: the same email address becomes the same token in every table, so foreign-key joins survive the mask and the drill still validates a database that answers production’s queries.

The tool never reverses a mask and never stores a mapping. It reads a masking policy, applies a keyed one-way transform whose key lives in a secret store, and proves its own work with a scan that runs as a separate pass. An unhandled sensitive column, a residual match, or a missing key each halts the drill rather than degrading quietly, because the failure mode of masking is data exposure.

Architecture and Execution Model

Figure. The tool loads its policy, applies the keyed transform per column, streams batched UPDATEs so large tables stay within budget, verifies its own output with an independent scan, and returns 0, 1, or 2 to gate the drill.

The transform is a keyed HMAC-SHA256 over the plaintext, truncated and shaped to preserve the format the schema and application depend on. Because HMAC is a pure function of the value and the session key, the same input yields the same token everywhere, which is exactly what keeps orders.customer_id joined to customers.id after both are masked. The scan is deliberately independent of the masker: it re-reads every column in each table and asserts that no PII pattern matches, except where a column’s own masking strategy legitimately preserves that shape, so a defect in the transform — or an entirely unhandled column the schema grew — cannot slip past the check meant to catch it. The whole run is idempotent at table granularity — a re-run over an already-masked table re-tokenizes already-tokenized strings, which stay non-identifying — so a crash mid-drill is recovered by re-running, not by restoring again.

Prerequisites

Python 3.9+ (the tool uses hmac, hashlib, and standard typing; dict-based policies parse from JSON).
The psycopg (v3) driver with its binary build, installed into the automation environment:
bash
```
pip install "psycopg[binary]>=3.1"
```
A freshly restored PostgreSQL sandbox reachable from the automation host, provisioned and network-isolated by the upstream drill phases. The tool connects with a role that can UPDATE the target tables.
A masking policy file (JSON) listing each table, its primary key, and the columns to mask with a per-column strategy (hmac, email, phone, or redact).
The masking key injected from a secret store into the DR_MASK_KEY environment variable at run time — never committed and never written to disk. The same key must be used for every column in a foreign-key relationship, or the join breaks.
A residual-PII pattern set and column allowlist so the scan knows what “clean” means: regexes for the shapes that must not survive, and the set of columns permitted to remain unmasked (surrogate ids, timestamps).

Production Implementation

The tool is structured as three stages behind one main: a masker that produces the keyed transform, a batched updater that streams large tables by primary-key ranges, and a scanner that independently verifies the result. It connects with psycopg, but the transform and scan logic are pure Python so the script runs and self-tests without a live database when invoked with --selftest.

python

#!/usr/bin/env python3
"""Deterministically mask PII in a freshly restored Postgres sandbox.

Applies a keyed HMAC transform to configured columns so foreign-key joins
survive, batches UPDATEs across large tables, then scans for residual PII.

Exit codes (consumed by the DR drill orchestrator):
    0  masking complete and residual scan clean -> release to validation
    1  residual PII found after masking          -> quarantine sandbox
    2  usage / policy / connection error         -> abort the drill
"""
from __future__ import annotations

import argparse
import hashlib
import hmac
import json
import os
import re
import sys
from dataclasses import dataclass
from typing import Callable, Dict, List, Optional

# Shapes the scan must not find on an unmasked column. The masked email uses
# the reserved .example TLD, so a masked address never matches the email rule.
PII_PATTERNS = {
    "email": re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.(?!example\b)[A-Za-z]{2,}"),
    "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
    "e164": re.compile(r"\+\d{11,15}\b"),
    "phone_us": re.compile(r"\(\d{3}\)\s?\d{3}-\d{4}"),
}

# Format-preserving strategies keep a shape on purpose, so a column masked that
# way is allow-listed against the pattern it legitimately still matches.
STRATEGY_ALLOWED = {"phone": frozenset({"e164", "phone_us"})}


@dataclass(frozen=True)
class TablePolicy:
    """Masking rules for one table."""

    table: str
    primary_key: str
    columns: Dict[str, str]  # column -> strategy


class Masker:
    """Deterministic keyed transforms; the key never leaves this process."""

    def __init__(self, key: bytes) -> None:
        if not key:
            raise ValueError("DR_MASK_KEY is empty; refuse to run unkeyed")
        self._key = key

    def _tok(self, value: str, n: int) -> str:
        return hmac.new(self._key, value.encode("utf-8"),
                        hashlib.sha256).hexdigest()[:n]

    def hmac(self, value: str) -> str:
        return "" if value == "" else self._tok(value, 16)

    def email(self, value: str) -> str:
        if "@" not in value:
            return self._tok(value, 12)
        local, _, domain = value.partition("@")
        # Trailing .example is a reserved TLD, so the scan's email pattern
        # (which excludes .example) cannot match the masked address.
        return f"{self._tok(local, 10)}@{self._tok(domain, 8)}.example"

    def phone(self, value: str) -> str:
        digits = [c for c in value if c.isdigit()]
        stream = self._tok(value, len(digits) or 1)
        out, i = [], 0
        for ch in value:
            if ch.isdigit():
                out.append(str(int(stream[i % len(stream)], 16) % 10))
                i += 1
            else:
                out.append(ch)
        return "".join(out)

    def redact(self, value: str) -> str:
        return "" if value == "" else "REDACTED"

    def resolve(self, strategy: str) -> Callable[[str], str]:
        fn = {"hmac": self.hmac, "email": self.email,
              "phone": self.phone, "redact": self.redact}.get(strategy)
        if fn is None:
            raise ValueError(f"unknown masking strategy: {strategy!r}")
        return fn


def load_policy(path: str) -> List[TablePolicy]:
    raw = json.loads(open(path, encoding="utf-8").read())
    return [TablePolicy(t["table"], t["primary_key"], dict(t["columns"]))
            for t in raw["tables"]]


def mask_table(conn, masker: Masker, pol: TablePolicy, batch: int) -> int:
    """Stream one table by primary-key ranges and UPDATE masked values.

    Batching by ascending primary key keeps memory bounded and avoids a
    single table-wide lock on very large tables. Returns rows updated.
    """
    cols = list(pol.columns)
    fns = {c: masker.resolve(s) for c, s in pol.columns.items()}
    updated, last_pk = 0, None
    select_cols = ", ".join([pol.primary_key, *cols])
    with conn.cursor() as read_cur, conn.cursor() as write_cur:
        while True:
            if last_pk is None:
                read_cur.execute(
                    f"SELECT {select_cols} FROM {pol.table} "
                    f"ORDER BY {pol.primary_key} LIMIT %s", (batch,))
            else:
                read_cur.execute(
                    f"SELECT {select_cols} FROM {pol.table} "
                    f"WHERE {pol.primary_key} > %s "
                    f"ORDER BY {pol.primary_key} LIMIT %s", (last_pk, batch))
            rows = read_cur.fetchall()
            if not rows:
                break
            set_clause = ", ".join(f"{c} = %s" for c in cols)
            for row in rows:
                pk, values = row[0], row[1:]
                masked = [fns[c](v) if v is not None else None
                          for c, v in zip(cols, values)]
                write_cur.execute(
                    f"UPDATE {pol.table} SET {set_clause} "
                    f"WHERE {pol.primary_key} = %s", (*masked, pk))
                updated += 1
                last_pk = pk
            conn.commit()
    return updated


def scan_row(table: str, names: List[str],
             strategies: List[Optional[str]], values) -> List[str]:
    """Pure PII scan of one row; allow-lists format-preserving columns."""
    hits: List[str] = []
    for name, strategy, value in zip(names, strategies, values):
        if value is None:
            continue
        allowed = STRATEGY_ALLOWED.get(strategy, frozenset())
        for pname, pattern in PII_PATTERNS.items():
            if pname in allowed:
                continue
            if pattern.search(str(value)):
                hits.append(f"{table}.{name} matched {pname}")
    return hits


def scan_table(conn, pol: TablePolicy) -> List[str]:
    """Scan every column of a table, not only the ones we masked.

    Scanning all columns is what catches an unhandled sensitive column the
    schema grew that no policy entry covers. Columns with a policy entry
    carry their strategy so format-preserving masks are not false-flagged.
    """
    hits: List[str] = []
    with conn.cursor() as cur:
        cur.execute(f"SELECT * FROM {pol.table}")
        names = [d.name for d in cur.description]
        strategies = [pol.columns.get(n) for n in names]
        for row in cur:
            hits.extend(scan_row(pol.table, names, strategies, row))
    return hits


def run(conn, masker: Masker, policy: List[TablePolicy], batch: int) -> int:
    total = 0
    for pol in policy:
        total += mask_table(conn, masker, pol, batch)
    residual: List[str] = []
    for pol in policy:
        residual.extend(scan_table(conn, pol))
    print(f"MASKED rows={total} tables={len(policy)}")
    if residual:
        for h in residual[:20]:
            print(f"RESIDUAL_PII {h}", file=sys.stderr)
        return 1
    print("SCAN clean; releasing sandbox to validation")
    return 0


def selftest() -> int:
    """Prove determinism and scan behaviour without a live database."""
    m = Masker(key=b"selftest-key")
    # Determinism: identical inputs map to identical tokens (joins survive).
    assert m.email("a@x.com") == m.email("a@x.com"), "email not deterministic"
    assert m.hmac("k") == m.hmac("k") and m.hmac("k") != m.hmac("j")
    # A masked email uses the reserved .example TLD, so it never trips the rule.
    assert not PII_PATTERNS["email"].search(m.email("alice@example.com"))
    # A format-preserving phone column is allow-listed, so the scan stays clean.
    masked_phone = m.phone("+1 (555) 867-5309")
    assert scan_row("t", ["mobile"], ["phone"], [masked_phone]) == []
    # An unmasked address in a column with no policy entry IS flagged.
    assert scan_row("t", ["leaked"], [None], ["real@corp.com"]), "scan missed PII"
    print("SELFTEST ok: transforms deterministic, scan flags only real PII")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description="Mask PII in a restored sandbox.")
    ap.add_argument("--policy", help="path to masking policy JSON")
    ap.add_argument("--dsn", help="psycopg connection string to the sandbox")
    ap.add_argument("--batch", type=int, default=5000, help="rows per batch")
    ap.add_argument("--selftest", action="store_true",
                    help="run offline determinism/scan checks and exit")
    args = ap.parse_args()

    if args.selftest:
        return selftest()
    if not args.policy or not args.dsn:
        print("usage: mask_sandbox.py --policy P --dsn DSN [--batch N]",
              file=sys.stderr)
        return 2

    key = os.environ.get("DR_MASK_KEY", "").encode("utf-8")
    try:
        masker = Masker(key)
        policy = load_policy(args.policy)
    except (OSError, ValueError, KeyError, json.JSONDecodeError) as exc:
        print(f"config error: {exc}", file=sys.stderr)
        return 2

    try:
        import psycopg  # imported lazily so --selftest needs no driver
    except ImportError:
        print("psycopg not installed: pip install 'psycopg[binary]'",
              file=sys.stderr)
        return 2

    try:
        with psycopg.connect(args.dsn) as conn:
            return run(conn, masker, policy, args.batch)
    except psycopg.Error as exc:
        print(f"database error: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    sys.exit(main())

The masker resolves a strategy per column, so one policy file drives HMAC tokenization for join keys, format-preserving transforms for parsed fields, and redaction for free text. The batched updater walks each table by ascending primary key with a WHERE pk > last cursor rather than OFFSET, so cost stays constant per batch on tables of any size and no single statement holds a table-wide lock. The scanner runs after every table is masked and reports the specific table.column and pattern for each surviving match, which is what turns a 1 exit into an actionable quarantine rather than an opaque failure.

A NULL value is passed through untouched rather than tokenized, because a masked NULL would violate the nullability the schema and application depend on and would also change query results the validation compares against. Empty strings are likewise preserved as empty, so a column’s cardinality of blank versus populated rows is unchanged — the transform hides identity, not structure. Every strategy shares the same _tok helper and the same session key, which is what guarantees that a value masked as an email in one table and referenced as a raw string elsewhere would still collide consistently if both columns were declared join partners; the policy, not the code, decides which columns share a transform, so extending coverage to a new table is a data change rather than a code change. That separation keeps the masker itself small enough to audit line by line, which matters for a component whose failure mode is disclosure.

Step-by-Step Execution Walkthrough

Self-test the transforms first. Run python3 mask_sandbox.py --selftest in CI to confirm the masks are deterministic and leave no residual PII before the tool ever touches a database.
Render the masking key from your secret store into DR_MASK_KEY for the run, using the same key that masked any table this drill will join against.

Point the tool at the restored sandbox, capturing its exit code:

bash

DR_MASK_KEY="$(vault kv get -field=key secret/dr/mask)" \
  python3 mask_sandbox.py --policy policy.json \
    --dsn "postgresql://masker@sandbox.internal:5432/appdb"; echo "exit=$?"

Branch on the exit code. 0 releases the sandbox to the smoke-test stage; 1 quarantines it because residual PII survived and must be investigated; 2 aborts on a bad policy, missing key, or connection failure.
Read the residual lines on a 1. Each RESIDUAL_PII line names the exact column and pattern, so you can add the missing column to the policy and re-run rather than re-restoring.
Re-run freely on interruption. The transform is idempotent, so a crashed run is recovered by re-invoking the tool; already-masked tables re-tokenize to the same non-identifying values.

Verification and Expected Output

A clean run masks every configured column, finds no residual match, and exits 0:

text

MASKED rows=4821004 tables=6
SCAN clean; releasing sandbox to validation
exit=0

A run that finds an un-masked value — typically a column the schema grew that the policy never covered — prints the offending columns and exits 1:

text

MASKED rows=4821004 tables=6
RESIDUAL_PII customers.backup_email matched email
exit=1

The exit code is the contract the orchestrator reads: 0 means the sandbox is safe to validate, 1 means identifiable data survived and the drill must stop, and 2 means the invocation or environment is broken and needs a fix before rerun.

Failure Modes and Troubleshooting

Symptom	Cause	Remediation
Exit `1`, `RESIDUAL_PII ... matched email`	A sensitive column exists with no policy entry	Add the column to `policy.json` with the correct strategy and re-run
Joins return no rows after masking	A foreign key and its target were masked with different keys or strategies	Mask every column in the relationship with the same strategy and the same `DR_MASK_KEY`
Exit `2`, `DR_MASK_KEY is empty`	The secret was not injected into the environment	Render the key from the secret store into `DR_MASK_KEY` before invoking the tool
Masking a large table times out or bloats WAL	Batch size too large, or autovacuum not keeping up with dead tuples	Lower `--batch`, and `VACUUM` the table between drills; the sandbox is disposable so aggressive settings are safe
Exit `2`, `database error`	Sandbox unreachable or the role lacks `UPDATE`	Confirm the sandbox DSN and that the masking role can update the target tables
Scan clean but distribution looks identifying	Low-cardinality column masked deterministically preserves frequencies	Switch that standalone column to a salted transform, accepting it can no longer be joined on

Because the scan runs as an independent pass, a clean 0 is positive evidence that no configured pattern survived — not merely that the masker claims to have run. Persist the masked row counts, the policy version, and the scan verdict to an immutable sink so every drill carries proof that identifiable data never left the production trust zone.

Integration Notes

The tool is built for headless orchestration: flags in, parseable lines out, strict POSIX codes. A thin wrapper turns the exit code into a release decision or a quarantine:

bash

#!/usr/bin/env bash
set -euo pipefail

DR_MASK_KEY="$(vault kv get -field=key secret/dr/mask)" \
  python3 mask_sandbox.py --policy "$POLICY" --dsn "$SANDBOX_DSN"
case $? in
  0) echo "masked" > /run/drill/mask_ok ;;
  1) curl -s -X POST "$PAGERDUTY_WEBHOOK" -d '{"event":"residual_pii"}'; exit 1 ;;
  *) echo "[$(date -u)] masker misconfigured; aborting"; exit 2 ;;
esac

Wire that wrapper into whichever scheduler owns the drill. In Airflow, invoke it from a BashOperator positioned after the restore task and before the smoke-test task, so a non-zero exit fails the DAG and short-circuits validation. In Celery, wrap the call in a task that raises on exit 1 so an event-driven drill records the residual-PII finding and dispatches escalation with low latency. Under cron or systemd, the strict exit codes let an OnFailure handler route the alert without extra glue. However it is scheduled, masking is the mandatory gate between restore and validation described in the parent guide, drill data masking and anonymization — its 0 is the only path to publishing the sandbox.

Frequently Asked Questions

Why HMAC rather than a plain hash or random substitution?

A plain hash is deterministic but unkeyed, so anyone who guesses a candidate value can confirm it by hashing it — a dictionary attack against low-entropy fields like emails or phone numbers. HMAC keys the transform with a secret from a secret store, so the mapping is stable and irreversible without the key. Random substitution would be unlinkable but would break every foreign-key join, which defeats the purpose of validating a real restore.

How does batching keep large-table masking within the drill budget?

The updater walks each table by ascending primary key with a WHERE pk > last cursor and commits per batch, so memory and lock scope stay bounded regardless of table size and cost per batch is constant. Using keyset pagination instead of OFFSET avoids the quadratic scan cost that OFFSET incurs on deep pages. Because the sandbox is disposable, you can also lower the batch size and vacuum aggressively without any production impact.

Why run the residual scan separately from the masker?

Because a check that shares code with the thing it checks can share the same bug. The scan re-reads the masked columns and applies an independent pattern set, so a defect in a transform that leaves data partially identifiable is still caught. A single surviving match returns exit 1 and quarantines the sandbox, which is the behavior a confidentiality gate must have rather than trusting the masker's own claim of success.

Is it safe to re-run the tool after a crash mid-drill?

Yes. Every transform is a pure function of the value and the session key, so re-masking an already-masked value re-tokenizes an already-tokenized string, which remains non-identifying. The tool is therefore idempotent at table granularity and a crashed run is recovered by re-invoking it rather than restoring the backup again. The final scan still runs and still gates the release, so a partial run can never be mistaken for a complete one.

Drill data masking and anonymization — the parent guide that defines masking as the mandatory post-restore, pre-validation gate this tool implements.
Sandbox provisioning — stands up and restores the isolated instance this tool connects to.
Security boundaries for DR environments — the network-isolation invariants masking complements as defense in depth.
Smoke-test routing logic — the validation stage that may only run once this tool reports a clean scan.

This script is one component of the broader drill data masking and anonymization guide.

For authoritative behavior of the primitives used here, consult the Python documentation for hmac and hashlib, and the PostgreSQL documentation on UPDATE.