Output Validation with Pydantic and Zod Against Malicious Content

from __future__ import annotations
 
from enum import Enum
from typing import Annotated, Literal
from urllib.parse import urlparse
 
from pydantic import BaseModel, Field, ValidationError, field_validator
 
 
class Priority(str, Enum):
    low = "low"
    normal = "normal"
    high = "high"
 
 
class CreateTicketArgs(BaseModel):
    title: Annotated[str, Field(min_length=3, max_length=120)]
    body: Annotated[str, Field(max_length=4000)]
    priority: Priority = Priority.normal
    # Block free-form "run this tool next" channels on a ticket object
    labels: list[Annotated[str, Field(max_length=32)]] = Field(default_factory=list, max_length=8)
 
    @field_validator("body")
    @classmethod
    def no_credential_dump_markers(cls, v: str) -> str:
        lowered = v.lower()
        for needle in ("begin private key", "api_key=", "authorization: bearer"):
            if needle in lowered:
                raise ValueError("body looks like a secret dump")
        return v
 
 
class HttpGetArgs(BaseModel):
    url: Annotated[str, Field(max_length=500)]
 
    @field_validator("url")
    @classmethod
    def only_https_allowlisted_hosts(cls, v: str) -> str:
        parsed = urlparse(v)
        if parsed.scheme != "https":
            raise ValueError("only https allowed")
        host = (parsed.hostname or "").lower()
        allowed = {"api.example.com", "status.example.com"}
        if host not in allowed:
            raise ValueError(f"host not allowlisted: {host}")
        return v
 
 
class AgentAction(BaseModel):
    type: Literal["create_ticket", "http_get", "reply"]
    create_ticket: CreateTicketArgs | None = None
    http_get: HttpGetArgs | None = None
    reply: Annotated[str, Field(max_length=2000)] | None = None
 
 
def dispatch(action: dict) -> str:
    try:
        act = AgentAction.model_validate(action)
    except ValidationError as e:
        return f"rejected: {e.error_count()} validation error(s)"
 
    if act.type == "create_ticket":
        if act.create_ticket is None:
            return "rejected: missing create_ticket payload"
        return f"ticket queued: {act.create_ticket.title!r}"
    if act.type == "http_get":
        if act.http_get is None:
            return "rejected: missing http_get payload"
        return f"would GET {act.http_get.url}"
    return f"reply: {(act.reply or '')[:80]}"
 
 
print(dispatch({
    "type": "http_get",
    "http_get": {"url": "https://evil.example/steal"},
}))
print(dispatch({
    "type": "create_ticket",
    "create_ticket": {
        "title": "Late package",
        "body": "Customer says package is late.",
        "priority": "high",
    },
}))

import { z } from "zod";
 
const CreateTicketArgs = z.object({
  title: z.string().min(3).max(120),
  body: z
    .string()
    .max(4000)
    .refine((v) => !/api_key=/i.test(v), { message: "secret-like body" }),
  priority: z.enum(["low", "normal", "high"]).default("normal"),
  labels: z.array(z.string().max(32)).max(8).default([]),
});
 
const HttpGetArgs = z.object({
  url: z
    .string()
    .max(500)
    .url()
    .refine((u) => u.startsWith("https://"), { message: "https only" })
    .refine((u) => {
      const host = new URL(u).hostname;
      return ["api.example.com", "status.example.com"].includes(host);
    }, { message: "host not allowlisted" }),
});
 
const AgentAction = z.discriminatedUnion("type", [
  z.object({ type: z.literal("create_ticket"), create_ticket: CreateTicketArgs }),
  z.object({ type: z.literal("http_get"), http_get: HttpGetArgs }),
  z.object({ type: z.literal("reply"), reply: z.string().max(2000) }),
]);
 
export function dispatch(action: unknown): string {
  const parsed = AgentAction.safeParse(action);
  if (!parsed.success) {
    return `rejected: ${parsed.error.issues.length} issue(s)`;
  }
  const act = parsed.data;
  if (act.type === "http_get") return `would GET ${act.http_get.url}`;
  if (act.type === "create_ticket") return `ticket queued: ${act.create_ticket.title}`;
  return `reply: ${act.reply.slice(0, 80)}`;
}

def run_tool(name: str, raw_args: dict, registry: dict) -> str:
    if name not in registry:
        return "rejected: unknown tool"
    schema = registry[name]
    try:
        args = schema.model_validate(raw_args)
    except ValidationError:
        return "rejected: invalid args"
    return registry[name].fn(args)  # only after validate