OpenAI (PII Redaction Guardrails)

pip install ironbook-sdk openai pydantic
python openai_iron.py

############################
# Defaults & entry points  #
############################

package policies.llm_guard

default allow := false

# Allow only when in scope AND there are NO deny reasons
allow if {
  in_scope
  not deny_nonempty
}

# Helper: true iff deny set has at least one element (binds x positively)
deny_nonempty if {
  count(deny) > 0
}

# Aggregate granular deny reasons (partial set in v1)
deny contains msg if { msg := reason_in_scope_not_matched }
deny contains msg if { msg := reason_capability_missing }
deny contains msg if { msg := reason_trust_low }
deny contains msg if { msg := reason_model_disallowed }
deny contains msg if { msg := reason_region_disallowed }
deny contains msg if { msg := reason_classification_disallowed }
deny contains msg if { msg := reason_purpose_disallowed }
deny contains msg if { msg := reason_pii_blocked }
deny contains msg if { msg := reason_secret_leak }
deny contains msg if { msg := reason_budget_missing }
deny contains msg if { msg := reason_budget_exceeded }
deny contains msg if { msg := reason_prompt_too_large }
deny contains msg if { msg := reason_attachment_disallowed }
deny contains msg if { msg := reason_url_limit_exceeded }
deny contains msg if { msg := reason_time_window_blocked }

############################
# Scope & configuration    #
############################

# Request scope
in_scope if {
  input.action == "infer"
  input.resource == "llm://responses"
}

# Tunables / allow-lists
allowed_models          = {"gpt-4o", "gpt-4o-mini", "o4-mini"}
allowed_classifications = {"public", "internal", "deidentified"}
allowed_regions         = {"US", "EU"}
allowed_purposes        = {"assistant", "summarization", "redaction"}

min_trust_threshold := 35
max_prompt_chars    := 20000
max_urls_in_prompt  := 10
max_attachments     := 3
require_cost_check  := true

# PII / secrets switches
block_on_any_pii  := false   # if true: any PII blocks; else: PII only allowed for purpose="redaction"
USE_EXTERNAL_SCAN := false   # if true: rely on input.context.pii_scan booleans instead of regexing prompt

############################
# Shortcuts / aliases      #
############################

prompt_str := sprintf("%v", [input.context.prompt])
model      := input.context.model
region     := input.context.region
class      := input.context.data_classification
purpose    := input.context.purpose

# Server-injected (do NOT pass from clients)
trust_score := input.trust
has_capability(cap) if { input.capabilities[_] == cap }

# Budget fields (may be undefined)
est_cost_cents := input.context.estimated_cost_cents
budget_cents   := input.context.daily_budget_remaining_cents

# Attachments: default 0 if missing
default attachments_count := 0
attachments_count := n if { n := count(input.context.attachments) }

# URL heuristics
contains_url if { regex.match(`(?i)https?://`, prompt_str) }
count_url_boolean_matches := n if {
  parts := split(lower(prompt_str), "http")
  n := count(parts) - 1
  n > 0
}
too_many_urls if {
  max_urls_in_prompt >= 0
  contains_url
  max_urls_in_prompt == 0
}
too_many_urls if {
  max_urls_in_prompt > 0
  contains_url
  count_url_boolean_matches > max_urls_in_prompt
}

############################
# PII / secrets detectors  #
############################

# External scan (privacy-friendly). Expect input.context.pii_scan booleans.
pii_flag_external if {
  USE_EXTERNAL_SCAN
  input.context.pii_scan[_] == true
}

# RE2 patterns (case-insensitive via (?i))
re_email          = `(?i)\b[0-9A-Z._%+-]+@[0-9A-Z.-]+\.[A-Z]{2,}\b`
re_phone          = `(?i)(?:\+?\d{1,3}[\s.-]?)?(?:\(?\d{2,4}\)?[\s.-]?)?\d{3}[\s.-]?\d{4}\b`
re_ssn_us         = `\b(?!000|666)\d{3}-(?!00)\d{2}-(?!0000)\d{4}\b`
re_ip4            = `\b(?:(?:25[0-5]|2[0-4]\d|1?\d?\d)\.){3}(?:25[0-5]|2[0-4]\d|1?\d?\d)\b`
re_ipv6           = `\b(?:[A-F0-9]{1,4}:){7}[A-F0-9]{1,4}\b`
re_credit_digits  = `\b(?:\d[ -]*?){13,19}\b`
re_iban           = `(?i)\b[A-Z]{2}\d{2}[A-Z0-9]{1,30}\b`
re_password_kv    = `(?i)\b(pass|password|passwd|pwd)\s*[:=]\s*['"]?[^'"\s]{6,}`
re_username_kv    = `(?i)\b(user(name)?|login)\s*[:=]\s*[^\s'"]{3,}`
re_aws_access_key = `\bAKIA[0-9A-Z]{16}\b`
re_gcp_key        = `(?i)\b[A-Za-z0-9_-]{20,}\.apps\.googleusercontent\.com\b`
re_gh_pat         = `\bghp_[A-Za-z0-9]{36}\b`
re_slack_token    = `\bxox[baprs]-[A-Za-z0-9-]{10,}\b`
re_private_key    = `-----BEGIN (?:RSA|DSA|EC|OPENSSH)?\s*PRIVATE KEY-----`
re_jwt            = `\b[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+\.[A-Za-z0-9-_]+\b`

# Internal regex path (only if not using external scan booleans)
pii_email    if { not USE_EXTERNAL_SCAN; regex.match(re_email,         prompt_str) }
pii_phone    if { not USE_EXTERNAL_SCAN; regex.match(re_phone,         prompt_str) }
pii_ssn      if { not USE_EXTERNAL_SCAN; regex.match(re_ssn_us,        prompt_str) }
pii_ip4      if { not USE_EXTERNAL_SCAN; regex.match(re_ip4,           prompt_str) }
pii_ip6      if { not USE_EXTERNAL_SCAN; regex.match(re_ipv6,          prompt_str) }
pii_credit   if { not USE_EXTERNAL_SCAN; regex.match(re_credit_digits, prompt_str) }
pii_iban     if { not USE_EXTERNAL_SCAN; regex.match(re_iban,          prompt_str) }
pii_username if { not USE_EXTERNAL_SCAN; regex.match(re_username_kv,   prompt_str) }
pii_password if { not USE_EXTERNAL_SCAN; regex.match(re_password_kv,   prompt_str) }

secret_aws     if { not USE_EXTERNAL_SCAN; regex.match(re_aws_access_key, prompt_str) }
secret_gcp     if { not USE_EXTERNAL_SCAN; regex.match(re_gcp_key,        prompt_str) }
secret_gh_pat  if { not USE_EXTERNAL_SCAN; regex.match(re_gh_pat,         prompt_str) }
secret_slack   if { not USE_EXTERNAL_SCAN; regex.match(re_slack_token,    prompt_str) }
secret_privkey if { not USE_EXTERNAL_SCAN; regex.match(re_private_key,    prompt_str) }
secret_jwt     if { not USE_EXTERNAL_SCAN; regex.match(re_jwt,            prompt_str) }

# "Any PII" = OR across detectors
any_pii if { pii_flag_external }
any_pii if { pii_email }
any_pii if { pii_phone }
any_pii if { pii_ssn }
any_pii if { pii_ip4 }
any_pii if { pii_ip6 }
any_pii if { pii_credit }
any_pii if { pii_iban }
any_pii if { pii_username }
any_pii if { pii_password }

# "Any secret" = OR across secret detectors
any_secret if { secret_aws }
any_secret if { secret_gcp }
any_secret if { secret_gh_pat }
any_secret if { secret_slack }
any_secret if { secret_privkey }
any_secret if { secret_jwt }

############################
# Deny reasons (strings)   #
############################

reason_in_scope_not_matched := msg if {
  not in_scope
  msg := "request out of scope for this policy"
}

reason_capability_missing := msg if {
  not has_capability("openai_infer")
  msg := "agent missing required capability openai_infer"
}

reason_trust_low := msg if {
  trust_score < min_trust_threshold
  msg := sprintf("trust score %v below threshold %v", [trust_score, min_trust_threshold])
}

reason_model_disallowed := msg if {
  not allowed_models[model]
  msg := sprintf("model %v not in allow-list", [model])
}

reason_region_disallowed := msg if {
  not allowed_regions[region]
  msg := sprintf("region %v not permitted", [region])
}

reason_classification_disallowed := msg if {
  not allowed_classifications[class]
  msg := sprintf("data classification %v not permitted for LLM usage", [class])
}

reason_purpose_disallowed := msg if {
  not allowed_purposes[purpose]
  msg := sprintf("purpose %v not permitted", [purpose])
}

# PII rules
reason_pii_blocked := msg if {
  block_on_any_pii
  any_pii
  msg := "PII detected in prompt; blocked by policy"
}
reason_pii_blocked := msg if {
  not block_on_any_pii
  any_pii
  purpose != "redaction"
  msg := "PII detected; only 'redaction' purpose is permitted"
}

reason_secret_leak := msg if {
  any_secret
  msg := "potential secret/credential detected in prompt"
}

reason_budget_missing := msg if {
  require_cost_check
  not valid_cost_fields
  msg := "missing cost/budget fields"
}
reason_budget_exceeded := msg if {
  require_cost_check
  valid_cost_fields
  to_number(est_cost_cents) > to_number(budget_cents)
  msg := sprintf("estimated cost %v¢ exceeds remaining daily budget %v¢", [est_cost_cents, budget_cents])
}
valid_cost_fields if {
  est_cost_cents != null
  budget_cents != null
}

reason_prompt_too_large := msg if {
  count(prompt_str) > max_prompt_chars
  msg := sprintf("prompt too large (%v chars > %v)", [count(prompt_str), max_prompt_chars])
}

reason_attachment_disallowed := msg if {
  max_attachments >= 0
  attachments_count > max_attachments
  msg := sprintf("too many attachments (%v > %v)", [attachments_count, max_attachments])
}

reason_url_limit_exceeded := msg if {
  too_many_urls
  msg := sprintf("too many URLs in prompt (limit %v)", [max_urls_in_prompt])
}

reason_time_window_blocked := msg if {
  input.context.change_window_enabled == true
  input.context.change_window_open != true
  msg := "LLM usage blocked outside approved change window"
}

"""
Interactive OpenAI + Iron Book Guardrails tester

- Requires Iron Book Python SDK >= v.0.3.2
- Accepts user input from the console
- Sends full `prompt` into policy context (so Rego can scan PII/secrets)
- Uses resource "llm://responses" to match the Rego v1 policy

Env vars:
  IRONBOOK_API_KEY: Get your key at https://ironbook.identitymachines.com
  OPENAI_API_KEY: Your Open AI API Key
  IRONBOOK_POLICY_ID: Get this by creating a policy (SDK or Iron Book UI
  IRONBOOK_AUDIENCE     (default: https://api.openai.com)
"""

import os
import re
import math
import asyncio
from typing import Dict, Any

from pydantic import BaseModel, Field

# Iron Book SDK (async)
from ironbook_sdk import (
    IronBookClient, GetAuthTokenOptions, PolicyInput, RegisterAgentOptions
)

# OpenAI SDK (sync)
from openai import OpenAI

# ------------------------------------------------------------------------------
# Configuration (env-overridable)
# ------------------------------------------------------------------------------
IRONBOOK_API_KEY = os.getenv('IRONBOOK_API_KEY', 'REPLACE ME')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY', 'REPLACE ME')

IRONBOOK_AUDIENCE = os.getenv('IRONBOOK_AUDIENCE', 'https://api.openai.com')
IRONBOOK_RESOURCE = 'llm://responses'  # matches the Rego v1 LLM Guard policy we’re testing

AGENT_NAME = os.getenv('IRONBOOK_AGENT_NAME', 'knowledgeops') #no dashes or special characters
AGENT_CAPABILITIES = os.getenv('IRONBOOK_AGENT_CAPABILITIES', '["openai_infer"]')
POLICY_ID = os.getenv('IRONBOOK_POLICY_ID', 'REPLACE ME')

# Simple per-model $/1K token pricing (illustrative)
MODEL_PRICING_PER_1K = {
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},   # $/1k tok (example)
    "gpt-4o":      {"input": 2.50, "output": 5.00},
    "o4-mini":     {"input": 0.30, "output": 1.20},
}

DEFAULT_DAILY_BUDGET_REMAINING_CENTS = 500  # $5.00 remaining

# ------------------------------------------------------------------------------
# Rough token & cost estimation
# ------------------------------------------------------------------------------
def estimate_tokens(text: str) -> int:
    return max(1, math.ceil(len(text) / 4))  # ~1 token ≈ 4 chars (very rough)

def estimate_cost_cents(model: str, prompt_tokens: int, max_output_tokens: int) -> int:
    prices = MODEL_PRICING_PER_1K.get(model, {"input": 5.0, "output": 10.0})
    input_cost = (prompt_tokens / 1000.0) * prices["input"]
    output_cost = (max_output_tokens / 1000.0) * prices["output"]
    total_usd = input_cost + output_cost
    return int(round(total_usd * 100))

# ------------------------------------------------------------------------------
# Iron Book authorize (async)
# ------------------------------------------------------------------------------
async def ib_authorize_openai_call(policy_id: str, context: Dict[str, Any]) -> Dict[str, Any]:
    iron = IronBookClient(api_key=IRONBOOK_API_KEY)

    # Register the agent in Iron Book if it doesn't exist
    # Note that trying to re-register an agent with the same name will return an error
    try:
        agent = await iron.register_agent(RegisterAgentOptions(
            agent_name=AGENT_NAME,
            capabilities=AGENT_CAPABILITIES
        ))
        print(f"\n🤖✅ AI agent registered with DID: {agent.did}\n")
    except Exception:
        # Fall back to getting an existing agent with that name if registration fails
        agent = await iron.get_agent(f"did:web:agents.identitymachines.com:{AGENT_NAME}")
        print(f"\n🤖✅ AI agent found with DID: {agent.did}\n")

    # Acquire a one-shot token for the OpenAI audience
    token_data = await iron.get_auth_token(GetAuthTokenOptions(
        agent_did=agent.did,
        vc=agent.vc,
        action="infer",
        resource=IRONBOOK_RESOURCE,
        audience=IRONBOOK_AUDIENCE
    ))
    access_token = token_data.get("access_token")
    if not access_token:
        raise RuntimeError("\n🪙❌ Failed to obtain Iron Book access token\n")

    print(f"\n🪙✅ AUTH success: Iron Book one-shot token issued to {AGENT_NAME} for action='infer' and resource='{IRONBOOK_RESOURCE}'\n")

    # Ask for the policy decision
    decision = await iron.policy_decision(PolicyInput(
        agent_did=agent.did,
        policy_id=policy_id,
        token=access_token,
        context=context
    ))

    # Normalize to dict (SDK may return object-like)
    if hasattr(decision, "dict"):
        try:
            decision = decision.dict()
        except Exception:
            decision = dict(decision)
    elif not isinstance(decision, dict):
        try:
            decision = dict(decision)
        except Exception:
            decision = {"allow": getattr(decision, "allow", False)}

    return decision

# ------------------------------------------------------------------------------
# Secured OpenAI call (sync wrapper)
# ------------------------------------------------------------------------------
class InferenceRequest(BaseModel):
    prompt: str = Field(..., description="User prompt text (will be sent to policy).")
    model: str = Field(default="gpt-4o-mini")
    max_output_tokens: int = Field(default=200)
    data_classification: str = Field(default="internal")  # 'public'|'internal'|'deidentified'
    purpose: str = Field(default="assistant")             # 'assistant'|'redaction'|...
    region: str = Field(default="US")
    daily_budget_remaining_cents: int = Field(default=DEFAULT_DAILY_BUDGET_REMAINING_CENTS)

def secured_openai_infer(req: InferenceRequest) -> str:
    # Rough cost estimate (used by policy budget check)
    prompt_tokens = estimate_tokens(req.prompt)
    est_cost_cents = estimate_cost_cents(req.model, prompt_tokens, req.max_output_tokens)

    # Build policy context (Rego v1 policy expects 'prompt' here)
    context = {
        "prompt": req.prompt,
        "model": req.model,
        "region": req.region,
        "data_classification": req.data_classification,
        "purpose": req.purpose,
        "estimated_cost_cents": est_cost_cents,
        "daily_budget_remaining_cents": req.daily_budget_remaining_cents,
        # Optional extras your policy might use later:
        # "attachments": [],
        # "change_window_enabled": False,
        # "change_window_open": True,
			# etc.
    }

    # Authorize via Iron Book (will not pass trust score or capabilities — those are server-injected)
    decision = asyncio.run(ib_authorize_openai_call(
        policy_id=POLICY_ID,
        context=context
    ))

    allowed = decision.get("allow", False)
    if not allowed:
        reason = (decision.get("reason")
                  or decision.get("message")
                  or decision.get("deny")
                  or decision)
        raise PermissionError(f"Denied by policy: {reason}")

    # If allowed, call OpenAI
    client = OpenAI(api_key=OPENAI_API_KEY)
    response = client.responses.create(
        model=req.model,
        input=req.prompt,
        max_output_tokens=req.max_output_tokens,
        store=False,
    )
    try:
        return getattr(response, "output_text", None) or str(response)
    except Exception:
        return str(response)

# ------------------------------------------------------------------------------
# CLI helpers
# ------------------------------------------------------------------------------
def read_multiline_prompt() -> str:
    print("\nEnter your prompt. Finish with a separate single line containing only 'END':")
    lines = []
    while True:
        try:
            line = input()
        except EOFError:
            break
        if line.strip() == "END":
            break
        lines.append(line)
    return "\n".join(lines).strip()

def ask_with_default(label: str, default: str) -> str:
    val = input(f"{label} [{default}]: ").strip()
    return val or default

def ask_int_with_default(label: str, default: int) -> int:
    val = input(f"{label} [{default}]: ").strip()
    if not val:
        return default
    try:
        return int(val)
    except ValueError:
        print("Not a valid integer; using default.")
        return default

# ------------------------------------------------------------------------------
# Main (interactive loop)
# ------------------------------------------------------------------------------
def main():
    # Basic env sanity
    missing = []
    if not IRONBOOK_API_KEY or IRONBOOK_API_KEY == 'REPLACE ME':
        missing.append("IRONBOOK_API_KEY")
    if not OPENAI_API_KEY or OPENAI_API_KEY == 'REPLACE ME':
        missing.append("OPENAI_API_KEY")
    if POLICY_ID.startswith('REPLACE'):
        print("⚠️  Set IRONBOOK_POLICY_ID env var to use your real policy.")
    if missing:
        print(f"⚠️  Missing env vars: {', '.join(missing)}")

    print("\nIron Book LLM Guard — interactive tester")
    print("Press Ctrl+C at any time to exit.")

    while True:
        try:
            prompt = read_multiline_prompt()
            if not prompt:
                print("No prompt entered. Exiting.")
                break

            model = ask_with_default("Model", "gpt-4o-mini")
            region = ask_with_default("Region", "US")
            data_classification = ask_with_default("Data classification (public|internal|deidentified)", "internal")
            purpose = ask_with_default("Purpose (assistant|summarization|redaction)", "assistant")
            max_output_tokens = ask_int_with_default("Max output tokens", 200)
            daily_budget_cents = ask_int_with_default("Daily budget remaining (cents)", DEFAULT_DAILY_BUDGET_REMAINING_CENTS)

            req = InferenceRequest(
                prompt=prompt,
                model=model,
                region=region,
                data_classification=data_classification,
                purpose=purpose,
                max_output_tokens=max_output_tokens,
                daily_budget_remaining_cents=daily_budget_cents
            )

            print("\n🔐 Authorizing via Iron Book policy...")
            try:
                output = secured_openai_infer(req)
                print("\n✅ ALLOW — OpenAI response:\n")
                print(output)
            except Exception as e:
                print(f"\n❌ DENY — {e}")

            again = input("\nRun another prompt? (y/N): ").strip().lower()
            if again != 'y':
                break

        except KeyboardInterrupt:
            print("\nExiting.")
            break

if __name__ == "__main__":
    main()