hncb-fusion-deid-demo/gateway_api/tokenize/handler.py

"""
/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 (small table, so a filtered scan is
    # fine). Best effort: if it fails we still tokenize, just without a RAG link.
    # NOTE: no `Limit` -- in DynamoDB, Limit caps items *scanned* before the filter
    # runs, so `Limit=1` returns nothing when the first row scanned isn't the match
    # (broke once a 2nd customer existed). A GSI on `name` would be the prod fix.
    if not CUSTOMERS:
        return None
    try:
        from boto3.dynamodb.conditions import Attr

        res = CUSTOMERS.scan(FilterExpression=Attr("name").eq(name))
        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)