Implement reversible PII de-identification round trip (T1–T7)

Build the AWS side of the HNCB demo end to end (region ap-southeast-1): - T1 /tokenize + T2 /restore: Lambdas behind a public API Gateway (shared-secret auth), Presidio detection, random per-request tokens, DynamoDB vault; overlap resolution so a ROC ID stays TW_ROC_ID. - T3: Presidio made private (SG-locked to the tokenize Lambda in-VPC; DynamoDB gateway endpoint); only /tokenize + /restore are public. - T4: RAG Lambda registered as an MCP tool on an AgentCore Gateway (AWS_IAM/SigV4); agentcore_setup.sh + a SigV4 MCP invoke test. - T5: Strands agent deployed to AgentCore Runtime; SigV4 gateway auth, apac inference profile, pinned deps. - T6: advisor UI on S3+CloudFront with a Fusion-less demo orchestrator (/demo) chaining tokenize -> runtime -> restore. - T7: README runbook + trace check; teardown deletes gateway/runtime/memory/ECR. Verified live: the cloud AgentCore/Bedrock trace shows only tokens, never the real name. Secrets stay in gitignored local.auto.tfvars. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-01 17:10:58 +10:00
parent 6b2a051be3
commit 78d67a2469
21 changed files with 1439 additions and 107 deletions
--- a/gateway_api/tokenize/handler.py
+++ b/gateway_api/tokenize/handler.py
@@ -0,0 +1,187 @@
+"""
+/tokenize endpoint (T1)  --  ingress de-identification for the HNCB demo (step 2).
+
+Fusion SaaS POSTs the advisor's raw query here. We:
+  1. send the text to the Presidio detector (/analyze) -> typed findings
+  2. mint a RANDOM token per finding (CUST_<rand> for a PERSON, TW_<rand> for a
+     Taiwan ROC ID) -- detection != redaction: the detector never sees a token,
+     WE own minting + the reversible map
+  3. write each mapping into the on-prem vault, keyed by a fresh session_id
+  4. splice the tokens back into the text (right-to-left, so offsets stay valid)
+  5. return {deidentified_prompt, session_id} -- the only thing that leaves for the cloud
+
+Vault row shape (kept compatible with lambda_rag/handler.py):
+    {token, type, value, session_id, expires_at [, customer_id]}
+  - `value` is the ORIGINAL pii string -> lets /restore (T2) put identity back
+  - `customer_id` is added for a PERSON we can resolve, so the RAG tool (T4)
+    can turn a token into a de-identified evidence package. The RAG handler reads
+    entry.get("customer_id") for non-"CUSTOMER" types, so this Just Works.
+
+Runs as a Lambda behind a Function URL (public HTTPS). Auth is a shared secret
+(Fusion sends it as a bearer token / x-api-key); see TOKENIZE_API_KEY below.
+Stdlib + boto3 only, per repo conventions.
+"""
+import base64
+import hmac
+import json
+import os
+import secrets
+import time
+import urllib.request
+import uuid
+
+import boto3
+
+ddb = boto3.resource("dynamodb")
+VAULT = ddb.Table(os.environ["VAULT_TABLE"])
+# Optional: resolving PERSON -> customer_id makes the vault row usable by the RAG
+# tool. If the customers table isn't configured we still tokenize (best effort).
+_CUSTOMERS_TABLE = os.environ.get("CUSTOMERS_TABLE")
+CUSTOMERS = ddb.Table(_CUSTOMERS_TABLE) if _CUSTOMERS_TABLE else None
+
+PRESIDIO_URL = os.environ.get("PRESIDIO_URL", "http://localhost:5001").rstrip("/")
+TTL_SECONDS = int(os.environ.get("VAULT_TTL_SECONDS", "3600"))
+# Shared secret Fusion presents. If unset we allow all callers (local dev only)
+# and say so in the response so it's obvious this isn't locked down.
+API_KEY = os.environ.get("TOKENIZE_API_KEY")
+
+# entity type -> token prefix. Only these are tokenized in the demo.
+_PREFIX = {"PERSON": "CUST", "TW_ROC_ID": "TW"}
+# When Presidio returns overlapping spans (e.g. the ROC ID digits flagged as BOTH
+# TW_ROC_ID and PERSON), the more specific entity wins. Higher = kept.
+_PRIORITY = {"TW_ROC_ID": 2, "PERSON": 1}
+
+
+def _resolve_overlaps(findings):
+    # Keep only demo-scoped entities, then greedily drop any finding that overlaps
+    # one we've already kept. Best-first order = specific entity, then score, then
+    # longer span -- so the ROC ID stays TW_ROC_ID and the spurious PERSON on the
+    # same digits is discarded (no double-tokenising, no corrupted splice).
+    cand = [f for f in findings if f["entity_type"] in _PREFIX]
+    cand.sort(
+        key=lambda f: (_PRIORITY.get(f["entity_type"], 0), f["score"], f["end"] - f["start"]),
+        reverse=True,
+    )
+    kept = []
+    for f in cand:
+        if any(f["start"] < k["end"] and f["end"] > k["start"] for k in kept):
+            continue  # overlaps a higher-priority finding already kept
+        kept.append(f)
+    return kept
+
+
+def _mint(prefix):
+    # Random per request -> the same query yields different tokens every call.
+    # 6 digits mirrors the seeded demo token (CUST_000123) for a familiar look.
+    return f"{prefix}_{secrets.randbelow(1_000_000):06d}"
+
+
+def _analyze(text, language="zh"):
+    body = json.dumps({"text": text, "language": language}).encode("utf-8")
+    req = urllib.request.Request(
+        f"{PRESIDIO_URL}/analyze",
+        data=body,
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+    with urllib.request.urlopen(req, timeout=10) as resp:
+        return json.loads(resp.read().decode("utf-8"))
+
+
+def _resolve_customer_id(name):
+    # Demo-grade name -> customer_id resolution (one seeded customer, so a scan is
+    # fine). Best effort: if it fails we still tokenize, just without a RAG link.
+    if not CUSTOMERS:
+        return None
+    try:
+        from boto3.dynamodb.conditions import Attr
+
+        res = CUSTOMERS.scan(FilterExpression=Attr("name").eq(name), Limit=1)
+        items = res.get("Items", [])
+        return items[0]["customer_id"] if items else None
+    except Exception:
+        return None
+
+
+def _authorized(headers):
+    if not API_KEY:
+        return True  # local dev / unset -> allow (flagged in the response)
+    presented = headers.get("x-api-key") or ""
+    if not presented:
+        auth = headers.get("authorization") or ""
+        if auth.lower().startswith("bearer "):
+            presented = auth[7:]
+    return bool(presented) and hmac.compare_digest(presented, API_KEY)
+
+
+def _parse(event):
+    """Accept a Lambda Function URL / API Gateway event, or a raw dict (tests)."""
+    if "body" in event and not isinstance(event.get("body"), dict):
+        raw = event.get("body") or "{}"
+        if event.get("isBase64Encoded"):
+            raw = base64.b64decode(raw).decode("utf-8")
+        payload = json.loads(raw) if raw else {}
+        headers = {k.lower(): v for k, v in (event.get("headers") or {}).items()}
+        return payload, headers
+    # Direct invoke: the event itself is the payload; no HTTP headers.
+    return event, {}
+
+
+def _reply(status, obj):
+    return {
+        "statusCode": status,
+        "headers": {"Content-Type": "application/json"},
+        "body": json.dumps(obj, ensure_ascii=False),
+    }
+
+
+def lambda_handler(event, _context):
+    payload, headers = _parse(event)
+
+    if not _authorized(headers):
+        return _reply(401, {"error": "unauthorized"})
+
+    query = (payload.get("query") or "").strip()
+    if not query:
+        return _reply(400, {"error": "query is required"})
+    language = payload.get("language", "zh")
+
+    try:
+        findings = _analyze(query, language)
+    except Exception as e:  # detector unreachable -> fail closed, never leak raw text
+        return _reply(502, {"error": f"detector unavailable: {e}"})
+
+    session_id = uuid.uuid4().hex
+    expires_at = int(time.time()) + TTL_SECONDS
+
+    # Drop overlapping/duplicate spans, then tokenize right-to-left so earlier
+    # offsets stay valid as we splice.
+    findings = _resolve_overlaps(findings)
+    findings.sort(key=lambda f: f["start"], reverse=True)
+    text = query
+    minted = []
+    for f in findings:
+        prefix = _PREFIX[f["entity_type"]]  # guaranteed present by _resolve_overlaps
+        original = f["text"]
+        token = _mint(prefix)
+
+        item = {
+            "token": token,
+            "type": f["entity_type"],
+            "value": original,          # original PII -> lets /restore put it back
+            "session_id": session_id,
+            "expires_at": expires_at,
+        }
+        if f["entity_type"] == "PERSON":
+            customer_id = _resolve_customer_id(original)
+            if customer_id:
+                item["customer_id"] = customer_id  # RAG tool resolves via this
+        VAULT.put_item(Item=item)
+
+        text = text[: f["start"]] + token + text[f["end"] :]
+        minted.append({"token": token, "type": f["entity_type"]})
+
+    resp = {"deidentified_prompt": text, "session_id": session_id, "tokens": minted}
+    if not API_KEY:
+        resp["warning"] = "TOKENIZE_API_KEY unset -- endpoint is unauthenticated (dev only)"
+    return _reply(200, resp)