Governed Diagnostic Data Slicing and PII Masking Broker for AI Agents

id: mcp-safe-diagnostic-data-broker
namespace: company.ai

description: |
  Expose a governed database diagnostic query and PII masking broker as an MCP tool for
  AI agents to inspect sanitized production-like data slices without exposing personal data.

inputs:
  - id: target_entity
    type: SELECT
    displayName: Target Entity Table
    description: Allowlisted entity table to query for diagnostics.
    defaults: customers
    values:
      - customers
      - billing_transactions
      - audit_events

  - id: lookup_filter
    type: STRING
    displayName: Diagnostic Filter Clause
    description: Parameterized SQL WHERE expression evaluated against the target entity.
    defaults: "status = 'SUSPENDED' AND account_tier = 'ENTERPRISE'"

  - id: purpose_justification
    type: STRING
    displayName: Investigation Justification
    description: Ticket ID or regulatory justification for auditing the diagnostic
      slice request.
    defaults: "INC-9042: investigating automated account suspension trigger"

  - id: max_records
    type: INT
    displayName: Maximum Records Ceiling
    description: Hard ceiling on returned records to prevent bulk data exfiltration.
    defaults: 25

  - id: redaction_mode
    type: SELECT
    displayName: PII Redaction Strategy
    description: Technique used to mask sensitive attributes before returning data
      to the AI agent.
    defaults: DETERMINISTIC_HASH
    values:
      - DETERMINISTIC_HASH
      - PARTIAL_MASK

  - id: dry_run
    type: BOOL
    displayName: Dry Run Mode
    description: When true, verifies query safety and row counts without
      materializing full payloads.
    defaults: false

tasks:
  - id: sandbox_workspace
    type: io.kestra.plugin.core.flow.WorkingDirectory
    description: Isolated workspace for database execution and Python anonymization
      processing.
    tasks:
      - id: initialize_mock_database
        type: io.kestra.plugin.jdbc.duckdb.Queries
        description: Sets up diagnostic relational tables containing realistic
          operational records and PII.
        url: "jdbc:duckdb:diagnostic.duckdb"
        fetchType: NONE
        sql: |
          CREATE TABLE IF NOT EXISTS customers (
            customer_id VARCHAR PRIMARY KEY,
            full_name VARCHAR,
            email VARCHAR,
            phone_number VARCHAR,
            ip_address VARCHAR,
            account_tier VARCHAR,
            status VARCHAR,
            created_at TIMESTAMP
          );

          CREATE TABLE IF NOT EXISTS billing_transactions (
            transaction_id VARCHAR PRIMARY KEY,
            customer_id VARCHAR,
            card_last_four VARCHAR,
            billing_email VARCHAR,
            amount_cents INTEGER,
            currency VARCHAR,
            status VARCHAR,
            created_at TIMESTAMP
          );

          CREATE TABLE IF NOT EXISTS audit_events (
            event_id VARCHAR PRIMARY KEY,
            customer_id VARCHAR,
            actor_email VARCHAR,
            action VARCHAR,
            client_ip VARCHAR,
            status VARCHAR,
            event_time TIMESTAMP
          );

          DELETE FROM customers;
          DELETE FROM billing_transactions;
          DELETE FROM audit_events;

          INSERT INTO customers VALUES
            ('cust_101', 'Alice Henderson', 'alice.henderson@enterprise-acme.com', '+1-555-0142', '198.51.100.24', 'ENTERPRISE', 'SUSPENDED', TIMESTAMP '2026-10-01 10:00:00'),
            ('cust_102', 'Marcus Vance', 'marcus.vance@vance-logistics.io', '+1-555-0198', '203.0.113.88', 'ENTERPRISE', 'ACTIVE', TIMESTAMP '2026-10-02 11:30:00'),
            ('cust_103', 'Elena Rostova', 'elena.rostova@globex-telecom.eu', '+44-20-7946-0912', '192.0.2.145', 'ENTERPRISE', 'SUSPENDED', TIMESTAMP '2026-10-03 14:15:00'),
            ('cust_104', 'David Chen', 'david.chen@chen-trading.org', '+1-555-0163', '198.51.100.99', 'GROWTH', 'ACTIVE', TIMESTAMP '2026-10-04 09:45:00');

          INSERT INTO billing_transactions VALUES
            ('txn_8801', 'cust_101', '4242', 'billing@enterprise-acme.com', 450000, 'USD', 'FAILED', TIMESTAMP '2026-10-01 10:05:00'),
            ('txn_8802', 'cust_103', '5555', 'ap@globex-telecom.eu', 890000, 'EUR', 'FAILED', TIMESTAMP '2026-10-03 14:20:00'),
            ('txn_8803', 'cust_102', '1234', 'finance@vance-logistics.io', 120000, 'USD', 'COMPLETED', TIMESTAMP '2026-10-02 11:35:00');

          INSERT INTO audit_events VALUES
            ('evt_01', 'cust_101', 'security-bot@enterprise-acme.com', 'API_KEY_REVOCATION', '198.51.100.24', 'SUSPENDED', TIMESTAMP '2026-10-01 10:01:00'),
            ('evt_02', 'cust_103', 'admin@globex-telecom.eu', 'WEBHOOK_TIMEOUT', '192.0.2.145', 'SUSPENDED', TIMESTAMP '2026-10-03 14:16:00');

      - id: validate_sql_filter_guard
        type: io.kestra.plugin.scripts.python.Script
        description: Enforces SQL safety guardrails and prevents destructive keywords or
          injection sequences.
        script: |
          import re
          import sys

          filter_clause = """{{ inputs.lookup_filter }}"""
          entity = """{{ inputs.target_entity }}"""

          allowed_entities = {"customers", "billing_transactions", "audit_events"}
          if entity not in allowed_entities:
              raise ValueError(f"Entity '{entity}' is not permitted by governance policy.")

          disallowed_patterns = [
              r';',
              r'--',
              r'/\*',
              r'\bDROP\b',
              r'\bALTER\b',
              r'\bTRUNCATE\b',
              r'\bINSERT\b',
              r'\bUPDATE\b',
              r'\bDELETE\b',
              r'\bGRANT\b',
              r'\bREVOKE\b',
              r'\bUNION\b',
              r'\bINTO\s+OUTFILE\b',
              r'\bEXEC\b',
              r'\bATTACH\b',
              r'\bCOPY\b'
          ]

          for pattern in disallowed_patterns:
              if re.search(pattern, filter_clause, re.IGNORECASE):
                  raise ValueError(f"Security policy violation: detected prohibited keyword or sequence matching '{pattern}'.")

          print(f"Filter guard passed for entity '{entity}': {filter_clause}")

      - id: execute_diagnostic_query
        type: io.kestra.plugin.jdbc.duckdb.Query
        description: Queries diagnostic records with strict limit boundaries against the
          isolated database.
        url: "jdbc:duckdb:diagnostic.duckdb"
        fetchType: FETCH
        sql: "SELECT * FROM {{ inputs.target_entity }} WHERE {{ inputs.lookup_filter }}
          LIMIT {{ inputs.max_records }}"

      - id: anonymize_and_mask_records
        type: io.kestra.plugin.scripts.python.Script
        description: Hashes entity IDs and redacts PII before computing cryptographic
          audit receipt.
        inputFiles:
          raw_records.json: "{{ outputs.execute_diagnostic_query.rows | json }}"
        outputFiles:
          - sanitized_diagnostic_slice.json
        script: |
          import json
          import re
          import hashlib
          import hmac
          from kestra import Kestra

          with open("raw_records.json", "r", encoding="utf-8") as f:
              records = json.load(f)

          redaction_mode = "{{ inputs.redaction_mode }}"
          is_dry_run = {{ inputs.dry_run }}
          salt = "{{ secret('DATA_MASKING_SALT') ?? 'diagnostic-salt-prod-2026' }}"

          def hash_identifier(val: str, prefix: str) -> str:
              if not val:
                  return val
              digest = hmac.new(salt.encode("utf-8"), str(val).encode("utf-8"), hashlib.sha256).hexdigest()
              return f"{prefix}_{digest[:10]}"

          def mask_email(email: str) -> str:
              if not email or "@" not in email:
                  return email
              user, domain = email.split("@", 1)
              if redaction_mode == "DETERMINISTIC_HASH":
                  digest = hashlib.sha256((salt + email).encode("utf-8")).hexdigest()[:12]
                  return f"user_{digest}@{domain}"
              else:
                  masked_user = user[0] + "***" + (user[-1] if len(user) > 1 else "")
                  return f"{masked_user}@{domain}"

          def mask_phone(phone: str) -> str:
              if not phone:
                  return phone
              digits = re.sub(r'\D', '', phone)
              last4 = digits[-4:] if len(digits) >= 4 else "0000"
              return f"+X-XXX-XXX-{last4}"

          def mask_ip(ip: str) -> str:
              if not ip:
                  return ip
              parts = ip.split(".")
              if len(parts) == 4:
                  return f"{parts[0]}.{parts[1]}.***.***"
              return "xxx.xxx.xxx.xxx"

          sanitized = []
          for row in records:
              clean = dict(row)
              if "customer_id" in clean:
                  clean["customer_id"] = hash_identifier(clean["customer_id"], "cust")
              if "transaction_id" in clean:
                  clean["transaction_id"] = hash_identifier(clean["transaction_id"], "txn")
              if "event_id" in clean:
                  clean["event_id"] = hash_identifier(clean["event_id"], "evt")

              for key in list(clean.keys()):
                  if "email" in key.lower():
                      clean[key] = mask_email(clean[key])
                  elif "phone" in key.lower():
                      clean[key] = mask_phone(clean[key])
                  elif "ip" in key.lower():
                      clean[key] = mask_ip(clean[key])
                  elif "card" in key.lower():
                      clean[key] = "****"
                  elif key.lower() == "full_name":
                      name = clean[key]
                      clean[key] = name[0] + ". [REDACTED]" if name else "[REDACTED]"

              sanitized.append(clean)

          output_payload = {
              "target_entity": "{{ inputs.target_entity }}",
              "purpose_justification": "{{ inputs.purpose_justification }}",
              "record_count": len(sanitized),
              "redaction_mode": redaction_mode,
              "dry_run": is_dry_run,
              "records": [] if is_dry_run else sanitized
          }

          serialized = json.dumps(output_payload, indent=2, default=str)
          receipt_digest = hashlib.sha256(serialized.encode("utf-8")).hexdigest()

          with open("sanitized_diagnostic_slice.json", "w", encoding="utf-8") as f:
              f.write(serialized)

          Kestra.outputs({
              "record_count": len(sanitized),
              "sha256_audit_receipt": receipt_digest,
              "preview": sanitized[:5] if not is_dry_run else []
          })

  - id: format_mcp_response
    type: io.kestra.plugin.core.output.OutputValues
    description: Packages sanitized diagnostics and cryptographic verification
      receipt for the MCP agent.
    values:
      status: SUCCESS
      record_count: "{{ outputs.anonymize_and_mask_records.vars.record_count ?? 0 }}"
      sha256_audit_receipt: "{{
        outputs.anonymize_and_mask_records.vars.sha256_audit_receipt ?? '' }}"
      preview: "{{ outputs.anonymize_and_mask_records.vars.preview ?? [] }}"
      storage_uri: "{{
        outputs.anonymize_and_mask_records.outputFiles['sanitized_diagnostic_sl\
        ice.json'] }}"

  - id: record_compliance_audit_log
    type: io.kestra.plugin.core.log.Log
    description: Generates an immutable execution audit record satisfying regulatory
      compliance standards.
    message: >
      [AUDIT MCP-DATA-BROKER] Request completed: entity={{ inputs.target_entity
      }}, justification="{{ inputs.purpose_justification }}", records={{
      outputs.format_mcp_response.values.record_count }}, sha256={{
      outputs.format_mcp_response.values.sha256_audit_receipt }}

outputs:
  - id: status
    type: STRING
    displayName: Broker Status
    description: Status indicator returned to the AI agent.
    value: "{{ outputs.format_mcp_response.values.status ?? 'FAILED' }}"

  - id: record_count
    type: INT
    displayName: Retrieved Record Count
    description: Number of diagnostic records retrieved and masked.
    value: "{{ outputs.format_mcp_response.values.record_count ?? 0 }}"

  - id: sha256_audit_receipt
    type: STRING
    displayName: Cryptographic Audit Receipt
    description: SHA-256 fingerprint generated over the sanitized payload.
    value: "{{ outputs.format_mcp_response.values.sha256_audit_receipt ?? '' }}"

  - id: records_preview
    type: JSON
    displayName: Masked Records Preview
    description: Preview slice of masked records returned directly to the agent.
    value: "{{ outputs.format_mcp_response.values.preview ?? [] }}"

  - id: storage_uri
    type: STRING
    displayName: Sanitized Artifact URI
    description: Internal storage location of the full redacted diagnostic dataset.
    value: "{{ outputs.format_mcp_response.values.storage_uri ?? '' }}"

triggers:
  - id: mcp_diagnostic_broker
    type: io.kestra.plugin.core.trigger.McpToolTrigger
    toolName: fetch_sanitized_diagnostic_slice
    title: Governed Diagnostic Data Slicing and PII Masking Broker
    toolDescription: >
      Extracts, validates, and masks diagnostic records from relational database
      tables. Call this tool when investigating customer errors, troubleshooting
      transaction failures, or debugging account states. Sensitive PII such as
      email addresses, phone numbers, IP addresses, and credit cards are
      automatically redacted or deterministically pseudonymized.
    mcpServer: default
    annotations:
      readOnly: true
      destructive: false
      idempotent: true
      openWorld: false
      returnDirect: false

errors:
  - id: log_policy_breach
    type: io.kestra.plugin.core.log.Log
    description: Logs rejected MCP query attempt with reason and security context.
    level: ERROR
    message: >
      [SECURITY MCP-DATA-BROKER] Request rejected by policy: entity={{
      inputs.target_entity }}, filter="{{ inputs.lookup_filter }}",
      justification="{{ inputs.purpose_justification }}"

  - id: format_error_mcp_response
    type: io.kestra.plugin.core.output.OutputValues
    description: Emits formatted rejection envelope to inform the calling AI agent
      of policy boundaries.
    values:
      status: REJECTED_BY_POLICY
      error_message: "The requested diagnostic query was rejected due to SQL safety or
        data governance constraints."