Build an Agent Security Framework

import { Redis } from "ioredis";
import { pool } from "../db";
import { randomBytes, createHash } from "node:crypto";
const redis = new Redis(process.env.REDIS_URL!);

interface SecurityPolicy {
  agentId: string;
  allowedTools: string[];
  deniedTools: string[];
  maxToolCallsPerTurn: number;
  requireApproval: string[];  // tools needing human approval
  sensitivePatterns: RegExp[];  // patterns to block in outputs
  maxTokenBudget: number;
}

interface SecurityCheck {
  passed: boolean;
  checks: Array<{ name: string; passed: boolean; details: string; severity: "critical" | "warning" | "info" }>;
  blocked: boolean;
  reason?: string;
}

const INJECTION_PATTERNS: Array<{ pattern: RegExp; severity: "critical" | "warning"; description: string }> = [
  { pattern: /ignore\s+(previous|above|all)\s+(instructions?|prompts?|rules?)/i, severity: "critical", description: "Direct instruction override attempt" },
  { pattern: /you\s+are\s+now\s+/i, severity: "critical", description: "Role reassignment attempt" },
  { pattern: /system\s*:\s*/i, severity: "warning", description: "System prompt injection" },
  { pattern: /\[\s*INST\s*\]/i, severity: "critical", description: "Instruction tag injection" },
  { pattern: /do\s+not\s+follow\s+(your|the|any)\s+(rules?|guidelines?|instructions?)/i, severity: "critical", description: "Rule bypass attempt" },
  { pattern: /pretend\s+(you|that|to\s+be)/i, severity: "warning", description: "Role-play injection" },
  { pattern: /\bDROP\s+TABLE\b|\bDELETE\s+FROM\b|\bTRUNCATE\b/i, severity: "critical", description: "SQL injection in tool args" },
  { pattern: /\brm\s+-rf\b|\bsudo\b|\bchmod\s+777\b/i, severity: "critical", description: "Shell injection in tool args" },
];

const SENSITIVE_OUTPUT_PATTERNS = [
  /\b\d{3}-\d{2}-\d{4}\b/,  // SSN
  /\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b/,  // Credit card
  /-----BEGIN.*PRIVATE KEY-----/,  // Private keys
  /(?:password|secret|token)\s*[=:]\s*[^\s]{8,}/i,  // Credentials
];

// Check input for prompt injection
export async function checkInput(agentId: string, input: string): Promise<SecurityCheck> {
  const checks: SecurityCheck["checks"] = [];
  let blocked = false;

  for (const pattern of INJECTION_PATTERNS) {
    if (pattern.pattern.test(input)) {
      checks.push({ name: "prompt_injection", passed: false, details: pattern.description, severity: pattern.severity });
      if (pattern.severity === "critical") blocked = true;
    }
  }

  if (checks.length === 0) checks.push({ name: "prompt_injection", passed: true, details: "No injection patterns detected", severity: "info" });

  // Log security event
  if (blocked) {
    await pool.query(
      `INSERT INTO agent_security_log (agent_id, event_type, details, severity, created_at) VALUES ($1, 'injection_blocked', $2, 'critical', NOW())`,
      [agentId, JSON.stringify(checks.filter((c) => !c.passed))]
    );
    await redis.hincrby("agent:security:stats", "injections_blocked", 1);
  }

  return { passed: !blocked, checks, blocked, reason: blocked ? checks.find((c) => !c.passed)?.details : undefined };
}

// Validate tool call before execution
export async function validateToolCall(agentId: string, toolName: string, args: Record<string, any>, policy: SecurityPolicy): Promise<SecurityCheck> {
  const checks: SecurityCheck["checks"] = [];
  let blocked = false;

  // Check if tool is allowed
  if (policy.deniedTools.includes(toolName)) {
    checks.push({ name: "tool_denied", passed: false, details: `Tool '${toolName}' is denied for this agent`, severity: "critical" });
    blocked = true;
  } else if (policy.allowedTools.length > 0 && !policy.allowedTools.includes(toolName)) {
    checks.push({ name: "tool_not_allowed", passed: false, details: `Tool '${toolName}' not in allowed list`, severity: "critical" });
    blocked = true;
  } else {
    checks.push({ name: "tool_permission", passed: true, details: `Tool '${toolName}' is allowed`, severity: "info" });
  }

  // Check tool call rate limit
  const callKey = `agent:calls:${agentId}:${Math.floor(Date.now() / 60000)}`;
  const callCount = await redis.incr(callKey);
  await redis.expire(callKey, 120);
  if (callCount > policy.maxToolCallsPerTurn) {
    checks.push({ name: "rate_limit", passed: false, details: `Exceeded ${policy.maxToolCallsPerTurn} tool calls/minute`, severity: "warning" });
    blocked = true;
  }

  // Check args for injection
  const argsStr = JSON.stringify(args);
  for (const pattern of INJECTION_PATTERNS) {
    if (pattern.pattern.test(argsStr)) {
      checks.push({ name: "args_injection", passed: false, details: `${pattern.description} in tool arguments`, severity: "critical" });
      blocked = true;
    }
  }

  // Check if tool requires human approval
  if (policy.requireApproval.includes(toolName)) {
    checks.push({ name: "requires_approval", passed: false, details: `Tool '${toolName}' requires human approval`, severity: "warning" });
    // Don't block — queue for approval
    await redis.rpush("agent:approval:queue", JSON.stringify({ agentId, toolName, args, requestedAt: new Date().toISOString() }));
  }

  // Audit log
  await pool.query(
    `INSERT INTO agent_security_log (agent_id, event_type, details, severity, created_at) VALUES ($1, 'tool_call', $2, $3, NOW())`,
    [agentId, JSON.stringify({ tool: toolName, args: argsStr.slice(0, 500), blocked }), blocked ? "warning" : "info"]
  );

  return { passed: !blocked, checks, blocked };
}

// Sanitize agent output before sending to user
export function sanitizeOutput(output: string): { sanitized: string; redacted: number } {
  let sanitized = output;
  let redacted = 0;

  for (const pattern of SENSITIVE_OUTPUT_PATTERNS) {
    const matches = sanitized.match(new RegExp(pattern.source, "g"));
    if (matches) {
      redacted += matches.length;
      sanitized = sanitized.replace(new RegExp(pattern.source, "g"), "[REDACTED]");
    }
  }

  return { sanitized, redacted };
}

// Get security dashboard
export async function getSecurityDashboard(agentId?: string): Promise<{
  totalEvents: number;
  injectionsBlocked: number;
  toolCallsBlocked: number;
  outputsRedacted: number;
  recentEvents: any[];
}> {
  const stats = await redis.hgetall("agent:security:stats");
  let sql = "SELECT * FROM agent_security_log";
  const params: any[] = [];
  if (agentId) { sql += " WHERE agent_id = $1"; params.push(agentId); }
  sql += " ORDER BY created_at DESC LIMIT 50";
  const { rows } = await pool.query(sql, params);

  return {
    totalEvents: rows.length,
    injectionsBlocked: parseInt(stats.injections_blocked || "0"),
    toolCallsBlocked: parseInt(stats.tools_blocked || "0"),
    outputsRedacted: parseInt(stats.outputs_redacted || "0"),
    recentEvents: rows,
  };
}