shannon/apps/worker/src/mcp-server/exploit-collector.ts

// Copyright (C) 2025 Keygraph, Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License version 3
// as published by the Free Software Foundation.

/**
 * Exploit Collector MCP Server (factory parameterized by vulnerability class
 * and per-run valid-ID set).
 *
 * Exposes a single Zod-validated MCP tool `add_exploit`, called once per
 * processed vulnerability by the 5 exploit-* agents (injection, xss, auth,
 * ssrf, authz). After the agent terminates, the host harvests
 * collector.getAll() and runs exploit-renderer to produce
 * {class}_exploitation_evidence.md. The collector state is the structured
 * output.
 *
 * Schema shape:
 *   - The SDK tool() helper consumes a ZodRawShape (flat object), not a
 *     top-level discriminated union. The visible shape is therefore a single
 *     z.object with common fields required, status as a string enum, and
 *     per-status fields marked optional at the SDK layer. Each field's
 *     `.describe()` text explains when it applies.
 *   - True per-status field enforcement runs inside the tool handler via a
 *     z.discriminatedUnion('status', ...). Missing-field errors come back to
 *     the agent as structured Zod issues with retryable=true so it can fix
 *     and retry the call.
 *
 * Strict queue-ID validation: vulnerability_id is refined against the per-run
 * queue's known IDs at schema-build time. Hallucinated or typo'd IDs are
 * rejected with a structured Zod error that includes the valid-ID list,
 * letting the agent recover locally.
 *
 * Each Zod schema's field-level descriptions carry the bullet labels and
 * reproducibility guidance, so the SDK injects it into the agent's tool
 * catalog.
 */

import type { McpSdkServerConfigWithInstance } from '@anthropic-ai/claude-agent-sdk';
import { createSdkMcpServer, tool } from '@anthropic-ai/claude-agent-sdk';
import { z } from 'zod';

// ============================================================================
// CLASS DISCRIMINATOR
// ============================================================================

export const EXPLOIT_VULN_CLASSES = ['injection', 'xss', 'auth', 'ssrf', 'authz'] as const;
export type VulnClass = (typeof EXPLOIT_VULN_CLASSES)[number];

// ============================================================================
// SCHEMA CONSTANTS
// ============================================================================

const SEVERITY_VALUES = ['critical', 'high', 'medium', 'low'] as const;
const CONFIDENCE_VALUES = ['high', 'medium', 'low'] as const;

const VALID_IDS_PREVIEW_LIMIT = 8;

function formatValidIdsPreview(validIds: ReadonlySet<string>): string {
  const list = [...validIds];
  const head = list.slice(0, VALID_IDS_PREVIEW_LIMIT).join(', ');
  return list.length > VALID_IDS_PREVIEW_LIMIT ? `${head}, … (${list.length} total)` : head;
}

// ============================================================================
// PUBLIC TYPES (discriminated union — what consumers see)
// ============================================================================

export type ExploitedExploit = {
  status: 'exploited';
  vulnerability_id: string;
  title: string;
  vulnerable_location: string;
  overview: string;
  prerequisites?: string | null;
  severity: (typeof SEVERITY_VALUES)[number];
  impact: string;
  exploitation_steps: string[];
  proof_of_impact: string;
  notes?: string | null;
};

export type BlockedExploit = {
  status: 'blocked';
  vulnerability_id: string;
  title: string;
  vulnerable_location: string;
  prerequisites?: string | null;
  confidence: (typeof CONFIDENCE_VALUES)[number];
  current_blocker: string;
  potential_impact: string;
  evidence_of_vulnerability: string;
  what_we_tried: string;
  how_this_would_be_exploited: string[];
  expected_impact: string;
  notes?: string | null;
};

export type AddExploitInput = ExploitedExploit | BlockedExploit;

// ============================================================================
// SCHEMA BUILDER
// ============================================================================

function buildSchemas(validIds: ReadonlySet<string>) {
  const vulnerabilityIdField = z
    .string()
    .min(1)
    .describe(
      'Vulnerability identifier (e.g. "INJ-VULN-03"). Must match an ID from this run\'s ' +
        '{class}_exploitation_queue.json exactly — the collector rejects IDs not in the queue. ' +
        `Valid IDs for this run: ${formatValidIdsPreview(validIds)}.`,
    )
    .refine((id: string) => validIds.has(id), {
      message:
        `Vulnerability ID not in this run's queue. Valid IDs: ` +
        `${formatValidIdsPreview(validIds)}. ` +
        'Check the queue.json for the canonical ID — likely a typo or hallucinated ID.',
    });

  const titleField = z
    .string()
    .min(1)
    .describe(
      'Descriptive vulnerability title (e.g. "SQL Injection — User Search", "IDOR — Unauthorized ' +
        'Access to User Orders"). Concise; encodes the vulnerability category and where it lives.',
    );

  const vulnerableLocationField = z
    .string()
    .min(1)
    .describe(
      'Endpoint or mechanism where the vulnerability exists (e.g. "GET /api/products?id=", ' +
        '"POST /login", or a code location like "controllers/userController.js:42").',
    );

  const overviewField = z
    .string()
    .min(1)
    .describe(
      'Brief summary of the exploit itself — what the vulnerability is and how it was demonstrated ' +
        '(or how it would be demonstrated, for blocked findings). 1-3 sentences.',
    );

  const prerequisitesField = z
    .string()
    .nullable()
    .optional()
    .describe(
      'Required setup, tools, or conditions to reproduce the exploit (e.g. authentication, ' +
        'specific role, prior application state). Omit or pass null when no prerequisites apply.',
    );

  const notesField = z
    .string()
    .nullable()
    .optional()
    .describe(
      'Optional supplementary context — caveats, related findings, environmental observations. ' +
        'Free-form Markdown. Omit or pass null when N/A.',
    );

  const statusField = z
    .enum(['exploited', 'blocked'])
    .describe(
      'Verdict bucket. Set to "exploited" only after reaching Proof of Exploitation Level 3+ with ' +
        'concrete impact evidence (extracted data, executed JavaScript, account takeover, internal ' +
        'service access). Set to "blocked" only for real vulnerabilities where external factors ' +
        '(NOT security defenses) prevented full exploitation. Findings where a security defense ' +
        'successfully prevented exploitation after exhaustive bypass attempts are FALSE POSITIVE — ' +
        'route those to your workspace tracking file, not this tool.',
    );

  // Per-status fields. All optional at the SDK shape layer because a single
  // ZodRawShape cannot express a top-level discriminated union; the handler
  // re-validates against the discriminated union below for true enforcement.
  const severityField = z
    .enum(SEVERITY_VALUES)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="exploited". Severity of the demonstrated impact. Critical = Level 4 ' +
        '(admin credentials extracted, sensitive data dumped, system commands executed, full account ' +
        'takeover). High = Level 3 (data extraction proven, authentication bypass confirmed, ' +
        'internal service access). Medium/Low based on impact narrowness or read-only access. Must ' +
        'reflect demonstrated impact, not theoretical potential.',
    );

  const impactField = z
    .string()
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="exploited". Business/security impact achieved by the exploit ' +
        '(e.g. "Extracted full user table including bcrypt password hashes for 1,247 users", ' +
        '"Achieved RCE as the application user; arbitrary shell commands executed"). Must describe ' +
        'what was actually demonstrated, not what could theoretically happen.',
    );

  const exploitationStepsField = z
    .array(z.string().min(1))
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="exploited". Ordered, reproducible exploitation steps — one Markdown ' +
        'blob per numbered step. Each step must include full URLs (protocol + domain + port + path ' +
        '+ params), complete payloads, and copy-paste-ready commands. Use clear placeholders for ' +
        'variable values like [SESSION_TOKEN], [DATABASE_NAME], [TABLE_NAME], [TARGET_USER_ID]. ' +
        'Write each step as natural Markdown — interleave prose with fenced code blocks (```bash, ' +
        '```http, etc.) as you would in a write-up. Steps must be detailed enough that someone ' +
        'unfamiliar with the application can follow without additional research.',
    );

  const proofOfImpactField = z
    .string()
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="exploited". Concrete evidence of successful exploitation — extracted ' +
        'data, achieved actions, captured request/response pairs, log excerpts. Markdown blob; ' +
        'interleave prose with fenced code blocks. Must show what the exploit demonstrably achieved, ' +
        'not theoretical impact.',
    );

  const confidenceField = z
    .enum(CONFIDENCE_VALUES)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="blocked". Confidence that this finding is a real vulnerability that ' +
        'would be exploited if the external blocker were removed. High = code analysis strongly ' +
        'confirms vulnerability and partial exploitation (Level 1-2) succeeded. Medium = code ' +
        'analysis confirms but live evidence is partial. Low = signal-only; revisit if blocker is ' +
        'removed in a future run.',
    );

  const currentBlockerField = z
    .string()
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="blocked". What prevents full exploitation (e.g. "Server crashes after ' +
        '5 requests, blocking enumeration", "OAuth callback requires verified third-party email ' +
        'account we could not provision"). Must be an external operational constraint, not a ' +
        'security defense.',
    );

  const potentialImpactField = z
    .string()
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="blocked". What could be achieved if the blocker were removed (e.g. ' +
        '"Full database read access", "Account takeover of arbitrary user via reset-token leak"). ' +
        'Distinct from impact — this is the hypothetical outcome, not a demonstrated one.',
    );

  const evidenceOfVulnerabilityField = z
    .string()
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="blocked". Code snippets, response excerpts, or observed behavior ' +
        'proving the vulnerability is real. Markdown blob; interleave prose with fenced code blocks. ' +
        'This is what convinces the reader the finding is not a false positive despite incomplete ' +
        'exploitation.',
    );

  const whatWeTriedField = z
    .string()
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="blocked". Log of attempted exploitation techniques and why each was ' +
        'blocked. Each attempt should document the payload, the observed result, and the inferred ' +
        'blocker. Markdown blob; multiple attempts as a list or distinct paragraphs. Demonstrates ' +
        'exhaustive bypass effort per the Bypass Exhaustion Protocol.',
    );

  const howThisWouldBeExploitedField = z
    .array(z.string().min(1))
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="blocked". Ordered hypothetical exploitation steps assuming the blocker ' +
        'is removed — one Markdown blob per numbered step. Same reproducibility requirements as ' +
        'exploitation_steps: full URLs, complete payloads, copy-paste-ready commands. Frame the ' +
        'first step as "If [blocker] were removed: …".',
    );

  const expectedImpactField = z
    .string()
    .min(1)
    .nullable()
    .optional()
    .describe(
      'REQUIRED when status="blocked". Specific data or access that would be compromised if ' +
        'exploitation succeeded (e.g. "Read access to all user profile data including PII; write ' +
        'access to user-owned resources"). Markdown blob.',
    );

  // The flat shape passed to tool(). The SDK uses this to build the agent's
  // tool catalog. Per-status enforcement happens in the handler via the
  // discriminated union below.
  const flatShape = {
    status: statusField,
    vulnerability_id: vulnerabilityIdField,
    title: titleField,
    vulnerable_location: vulnerableLocationField,
    overview: overviewField,
    prerequisites: prerequisitesField,
    notes: notesField,
    severity: severityField,
    impact: impactField,
    exploitation_steps: exploitationStepsField,
    proof_of_impact: proofOfImpactField,
    confidence: confidenceField,
    current_blocker: currentBlockerField,
    potential_impact: potentialImpactField,
    evidence_of_vulnerability: evidenceOfVulnerabilityField,
    what_we_tried: whatWeTriedField,
    how_this_would_be_exploited: howThisWouldBeExploitedField,
    expected_impact: expectedImpactField,
  };

  // Strict per-status validation. Re-runs in the handler so missing fields
  // for the chosen status return a retryable Zod error to the agent.
  const ExploitedSchema = z.object({
    status: z.literal('exploited'),
    vulnerability_id: vulnerabilityIdField,
    title: titleField,
    vulnerable_location: vulnerableLocationField,
    overview: overviewField,
    prerequisites: prerequisitesField,
    severity: z.enum(SEVERITY_VALUES),
    impact: z.string().min(1),
    exploitation_steps: z.array(z.string().min(1)).min(1),
    proof_of_impact: z.string().min(1),
    notes: notesField,
  });

  const BlockedSchema = z.object({
    status: z.literal('blocked'),
    vulnerability_id: vulnerabilityIdField,
    title: titleField,
    vulnerable_location: vulnerableLocationField,
    prerequisites: prerequisitesField,
    confidence: z.enum(CONFIDENCE_VALUES),
    current_blocker: z.string().min(1),
    potential_impact: z.string().min(1),
    evidence_of_vulnerability: z.string().min(1),
    what_we_tried: z.string().min(1),
    how_this_would_be_exploited: z.array(z.string().min(1)).min(1),
    expected_impact: z.string().min(1),
    notes: notesField,
  });

  const StrictSchema = z.discriminatedUnion('status', [ExploitedSchema, BlockedSchema]);

  return { flatShape, StrictSchema };
}

// ============================================================================
// RESPONSE HELPERS
// ============================================================================

interface ToolResult {
  [x: string]: unknown;
  content: Array<{ type: 'text'; text: string }>;
  isError: boolean;
}

function createToolResult(response: { status: string; [key: string]: unknown }): ToolResult {
  return {
    content: [{ type: 'text', text: JSON.stringify(response, null, 2) }],
    isError: response.status === 'error',
  };
}

function successResult(data: Record<string, unknown>): ToolResult {
  return createToolResult({ status: 'success', ...data });
}

function errorResult(message: string, errorType = 'ValidationError', retryable = true): ToolResult {
  return createToolResult({ status: 'error', message, errorType, retryable });
}

function formatZodIssues(error: z.ZodError): string {
  return error.issues
    .map((issue) => {
      const path = issue.path.length > 0 ? issue.path.join('.') : '(root)';
      return `- ${path}: ${issue.message}`;
    })
    .join('\n');
}

// ============================================================================
// SERVER FACTORY
// ============================================================================

export interface ExploitCollectorServer {
  server: McpSdkServerConfigWithInstance;
  getAll(): AddExploitInput[];
}

export interface CreateExploitCollectorOptions {
  vulnClass: VulnClass;
  validIds: ReadonlySet<string>;
}

export function createExploitCollector(options: CreateExploitCollectorOptions): ExploitCollectorServer {
  const { vulnClass, validIds } = options;
  const exploits: AddExploitInput[] = [];
  const { flatShape, StrictSchema } = buildSchemas(validIds);

  const addExploitTool = tool(
    'add_exploit',
    `Record a single processed ${vulnClass} vulnerability as structured exploitation evidence. ` +
      'Call this once per vulnerability in your queue.json after reaching a definitive verdict ' +
      '(either successfully exploited or potential-but-blocked). The status field discriminates the ' +
      "two report buckets; required sub-fields differ per status (see each field's description for " +
      'which status requires it). Duplicate vulnerability_id calls are rejected — each vuln may only ' +
      'be recorded once. Vulnerability IDs not in the queue.json are rejected with a list of valid ' +
      'IDs. FALSE POSITIVE findings do NOT use this tool — they go to your workspace tracking file. ' +
      'After all queue vulnerabilities have been emitted, the host renderer assembles the ' +
      'deliverable Markdown from your recorded calls.',
    flatShape,
    async (input): Promise<ToolResult> => {
      // Re-validate against the strict discriminated union for per-status enforcement.
      const parsed = StrictSchema.safeParse(input);
      if (!parsed.success) {
        return errorResult(
          `Schema validation failed for status="${(input as { status?: string }).status}". ` +
            'Required-field issues:\n' +
            formatZodIssues(parsed.error),
          'ValidationError',
          true,
        );
      }
      const typed = parsed.data as AddExploitInput;
      const existing = exploits.find((e) => e.vulnerability_id === typed.vulnerability_id);
      if (existing) {
        return errorResult(
          `Vulnerability ${typed.vulnerability_id} has already been recorded. Each vulnerability ` +
            'may only be added once. Reach a final verdict before emitting.',
          'DuplicateError',
          false,
        );
      }
      exploits.push(typed);
      return successResult({ added: [typed.vulnerability_id], recorded_status: typed.status });
    },
  );

  const server: McpSdkServerConfigWithInstance = createSdkMcpServer({
    name: 'exploit-collector',
    version: '1.0.0',
    tools: [addExploitTool],
  });

  return {
    server,
    getAll: (): AddExploitInput[] => [...exploits],
  };
}