ktg-plugin-marketplace/plugins/llm-security/scanners/lib/mcp-description-cache.mjs

// mcp-description-cache.mjs — Cache MCP tool descriptions and detect drift.
// Zero external dependencies.
//
// Purpose:
//   MCP servers can change tool descriptions between sessions (rug-pull, MCP05).
//   This module caches the first-seen description for each tool and alerts when
//   a subsequent invocation delivers a description that has drifted significantly
//   (Levenshtein distance > 10% of original length).
//
// v7.3.0 (E14): Adds a sticky baseline per tool so cumulative drift across
//   many small updates can be detected. Each entry now carries:
//     - description: latest description
//     - firstSeen / lastSeen: timestamps
//     - baseline: { description, seenAt } — immutable until clearBaseline()
//     - history: [{ description, seenAt, distance }] — last 10 drift events (FIFO)
//   Cumulative drift = levenshtein(current, baseline.description) / max(|current|, |baseline|)
//   When cumulative >= cumulative_drift_threshold (default 0.25), emit advisory.
//   Baseline survives TTL purge so slow-burn drift is preserved.
//
// Storage: ~/.cache/llm-security/mcp-descriptions.json
// TTL: 7 days per entry — but entries with a baseline survive purge.
//
// OWASP: MCP05 (Tool Description Manipulation / Rug Pull)

import { readFileSync, writeFileSync, mkdirSync, existsSync } from 'node:fs';
import { join, dirname } from 'node:path';
import { homedir } from 'node:os';
import { levenshtein } from './string-utils.mjs';
import { getPolicyValue } from './policy-loader.mjs';

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const CACHE_DIR = join(homedir(), '.cache', 'llm-security');
const CACHE_FILE = join(CACHE_DIR, 'mcp-descriptions.json');
const TTL_MS = 7 * 24 * 60 * 60 * 1000; // 7 days
const DRIFT_THRESHOLD = 0.10; // 10% Levenshtein distance per-update
const CUMULATIVE_DRIFT_THRESHOLD_DEFAULT = 0.25; // 25% baseline drift
const HISTORY_CAP = 10;

// ---------------------------------------------------------------------------
// Cache structure (v7.3.0)
// ---------------------------------------------------------------------------
// {
//   "mcp__server__tool": {
//     "description": "latest description text",
//     "firstSeen": 1712345678000,
//     "lastSeen": 1712345678000,
//     "hash": "sha256-prefix (optional, for quick equality check)",
//     "baseline": { "description": "...", "seenAt": 1712345678000 },
//     "history": [
//       { "description": "...", "seenAt": 1712345678000, "distance": 12 }
//     ]
//   }
// }

/**
 * Resolve cache file path. Env-var override useful for testing the hook
 * without polluting the user's real cache.
 * @param {object} opts
 * @returns {string}
 */
function resolveCacheFile(opts) {
  if (opts.cacheFile) return opts.cacheFile;
  if (process.env.LLM_SECURITY_MCP_CACHE_FILE) return process.env.LLM_SECURITY_MCP_CACHE_FILE;
  return CACHE_FILE;
}

/**
 * Migrate a legacy v7.2.0 cache entry to the v7.3.0 shape.
 * Discriminator: presence of `history` array. v7.2.0 entries have neither
 * `baseline` nor `history`; v7.3.0 entries always have `history` (even if
 * empty). After `clearBaseline()` an entry has `history` but no `baseline`,
 * which is NOT a legacy entry and must NOT be re-seeded here — the next
 * `checkDescriptionDrift` call will seed baseline from the incoming
 * description instead.
 * Idempotent — running on a v7.3.0 entry is a no-op.
 * @param {object} entry
 * @returns {object}
 */
function migrateEntry(entry) {
  if (!Array.isArray(entry.history)) {
    // v7.2.0 → v7.3.0: seed baseline from current description and add history
    if (!entry.baseline) {
      entry.baseline = {
        description: entry.description,
        seenAt: entry.firstSeen,
      };
    }
    entry.history = [];
  }
  return entry;
}

/**
 * Load the cache from disk. Purges entries older than TTL — but entries with
 * a `baseline` survive purge, so cumulative-drift detection persists across
 * the 7-day window.
 * Returns empty object if file doesn't exist or is corrupt.
 * Migrates v7.2.0 entries on the fly.
 * @param {object} [opts] - Options for testing
 * @param {string} [opts.cacheFile] - Override cache file path
 * @param {number} [opts.now] - Override current time
 * @returns {Record<string, object>}
 */
export function loadCache(opts = {}) {
  const cacheFile = resolveCacheFile(opts);
  const now = opts.now ?? Date.now();

  if (!existsSync(cacheFile)) return {};

  try {
    const raw = readFileSync(cacheFile, 'utf-8');
    const data = JSON.parse(raw);
    if (!data || typeof data !== 'object') return {};

    const cleaned = {};
    for (const [key, entry] of Object.entries(data)) {
      if (entry && typeof entry === 'object' && typeof entry.lastSeen === 'number') {
        const stale = now - entry.lastSeen > TTL_MS;
        if (!stale || entry.baseline) {
          cleaned[key] = migrateEntry(entry);
        }
      }
    }
    return cleaned;
  } catch {
    return {};
  }
}

/**
 * Save the cache to disk. Creates the cache directory if needed.
 * @param {Record<string, object>} cache
 * @param {object} [opts]
 * @param {string} [opts.cacheFile] - Override cache file path
 */
export function saveCache(cache, opts = {}) {
  const cacheFile = resolveCacheFile(opts);
  const dir = dirname(cacheFile);

  try {
    if (!existsSync(dir)) {
      mkdirSync(dir, { recursive: true });
    }
    writeFileSync(cacheFile, JSON.stringify(cache, null, 2), 'utf-8');
  } catch {
    // Silently fail — drift detection is advisory, not critical
  }
}

/**
 * Resolve the cumulative-drift threshold (0..1).
 * Order: opts.cumulativeThreshold → policy.json (mcp.cumulative_drift_threshold) → default.
 * @param {object} opts
 * @returns {number}
 */
function resolveCumulativeThreshold(opts) {
  if (typeof opts.cumulativeThreshold === 'number') return opts.cumulativeThreshold;
  try {
    const v = getPolicyValue('mcp', 'cumulative_drift_threshold', CUMULATIVE_DRIFT_THRESHOLD_DEFAULT);
    return typeof v === 'number' ? v : CUMULATIVE_DRIFT_THRESHOLD_DEFAULT;
  } catch {
    return CUMULATIVE_DRIFT_THRESHOLD_DEFAULT;
  }
}

/**
 * Check a tool description against the cached version.
 *
 * First call for a tool: caches the description AND seeds the baseline,
 * returns no drift.
 * Subsequent calls:
 *   - per-update drift: levenshtein(current, previous) / |previous| > 10%
 *   - cumulative drift: levenshtein(current, baseline) / max(|current|, |baseline|) >= 25%
 *
 * @param {string} toolName - Full tool name (e.g. "mcp__tavily__tavily_search")
 * @param {string} description - Current tool description
 * @param {object} [opts] - Options for testing
 * @param {string} [opts.cacheFile] - Override cache file path
 * @param {number} [opts.now] - Override current time
 * @param {number} [opts.cumulativeThreshold] - Override cumulative threshold (0..1)
 * @returns {{
 *   drift: boolean,
 *   detail: string|null,
 *   distance: number,
 *   threshold: number,
 *   cached: string|null,
 *   baselineDriftPct: number,
 *   perUpdateDriftPct: number,
 *   cumulative: { drifted: boolean, distance: number, threshold: number, detail: string|null, baseline: string|null }
 * }}
 */
export function checkDescriptionDrift(toolName, description, opts = {}) {
  const now = opts.now ?? Date.now();
  const noDrift = {
    drift: false,
    detail: null,
    distance: 0,
    threshold: 0,
    cached: null,
    baselineDriftPct: 0,
    perUpdateDriftPct: 0,
    cumulative: { drifted: false, distance: 0, threshold: 0, detail: null, baseline: null },
  };

  if (!toolName || !description || typeof description !== 'string') {
    return noDrift;
  }

  const cumulativeThreshold = resolveCumulativeThreshold(opts);
  const cache = loadCache(opts);
  const existing = cache[toolName];

  if (!existing) {
    // First time seeing this tool — cache it and seed the baseline
    cache[toolName] = {
      description,
      firstSeen: now,
      lastSeen: now,
      baseline: { description, seenAt: now },
      history: [],
    };
    saveCache(cache, opts);
    return {
      ...noDrift,
      threshold: DRIFT_THRESHOLD,
      cumulative: { drifted: false, distance: 0, threshold: cumulativeThreshold, detail: null, baseline: description },
    };
  }

  // Defensive — entries from loadCache are already migrated, but be paranoid
  migrateEntry(existing);

  // Baseline was explicitly cleared (clearBaseline) — re-seed from the
  // incoming description so the next legitimate MCP version becomes the
  // new baseline.
  if (!existing.baseline) {
    existing.description = description;
    existing.lastSeen = now;
    existing.baseline = { description, seenAt: now };
    saveCache(cache, opts);
    return {
      ...noDrift,
      threshold: DRIFT_THRESHOLD,
      cached: null,
      cumulative: { drifted: false, distance: 0, threshold: cumulativeThreshold, detail: null, baseline: description },
    };
  }

  // Update lastSeen
  existing.lastSeen = now;

  // Quick equality check
  if (existing.description === description) {
    saveCache(cache, opts);
    return {
      ...noDrift,
      threshold: DRIFT_THRESHOLD,
      cached: existing.description,
      cumulative: {
        drifted: false,
        distance: 0,
        threshold: cumulativeThreshold,
        detail: null,
        baseline: existing.baseline.description,
      },
    };
  }

  // Per-update Levenshtein distance (vs previous description)
  const perDist = levenshtein(existing.description, description);
  const perBaseLen = Math.max(existing.description.length, 1);
  const perRatio = perDist / perBaseLen;
  const perDrifted = perRatio > DRIFT_THRESHOLD;

  // Cumulative Levenshtein distance (vs baseline)
  const cumDist = levenshtein(existing.baseline.description, description);
  const cumDenom = Math.max(existing.baseline.description.length, description.length, 1);
  const cumRatio = cumDist / cumDenom;
  const cumDrifted = cumRatio >= cumulativeThreshold;

  // Push event into history (FIFO, capped)
  existing.history.push({ description, seenAt: now, distance: perDist });
  if (existing.history.length > HISTORY_CAP) {
    existing.history.splice(0, existing.history.length - HISTORY_CAP);
  }

  const cachedDesc = existing.description;
  // Update current description (baseline stays put)
  existing.description = description;
  saveCache(cache, opts);

  let perDetail = null;
  if (perDrifted) {
    const pct = Math.round(perRatio * 100);
    perDetail =
      `Tool "${toolName}" description changed by ${pct}% (${perDist} edits / ${perBaseLen} chars). ` +
      `Threshold: ${Math.round(DRIFT_THRESHOLD * 100)}%. This may indicate a rug-pull attack (OWASP MCP05).`;
  }

  let cumDetail = null;
  if (cumDrifted) {
    const pct = Math.round(cumRatio * 100);
    cumDetail =
      `Tool "${toolName}" cumulative description drift ${pct}% from baseline (${cumDist} edits). ` +
      `Threshold: ${Math.round(cumulativeThreshold * 100)}%. ` +
      `Slow-burn rug-pull may evade per-update detection (OWASP MCP05).`;
  }

  return {
    drift: perDrifted,
    detail: perDetail,
    distance: perDist,
    threshold: DRIFT_THRESHOLD,
    cached: cachedDesc,
    baselineDriftPct: cumRatio,
    perUpdateDriftPct: perRatio,
    cumulative: {
      drifted: cumDrifted,
      distance: cumDist,
      threshold: cumulativeThreshold,
      detail: cumDetail,
      baseline: existing.baseline.description,
    },
  };
}

/**
 * Extract MCP server name from a tool name.
 * Convention: mcp__<server>__<tool>
 * @param {string} toolName
 * @returns {string|null}
 */
export function extractMcpServer(toolName) {
  if (!toolName?.startsWith('mcp__')) return null;
  const parts = toolName.split('__');
  // mcp__server__tool → parts = ['mcp', 'server', 'tool']
  return parts.length >= 3 ? parts[1] : null;
}

/**
 * Clear the entire cache (for testing).
 * @param {object} [opts]
 * @param {string} [opts.cacheFile] - Override cache file path
 */
export function clearCache(opts = {}) {
  saveCache({}, opts);
}

/**
 * Clear the baseline slot for one tool (or all tools when toolName omitted).
 * Preserves description, firstSeen, lastSeen, and history. After clearing,
 * the next checkDescriptionDrift call will re-seed the baseline from the
 * current description.
 *
 * @param {string} [toolName] - Specific tool, or omit to clear all baselines
 * @param {object} [opts]
 * @param {string} [opts.cacheFile] - Override cache file path
 * @returns {{ cleared: number, tools: string[] }}
 */
export function clearBaseline(toolName, opts = {}) {
  const cache = loadCache(opts);
  const cleared = [];

  if (toolName) {
    if (cache[toolName] && cache[toolName].baseline) {
      delete cache[toolName].baseline;
      cleared.push(toolName);
    }
  } else {
    for (const [key, entry] of Object.entries(cache)) {
      if (entry && entry.baseline) {
        delete entry.baseline;
        cleared.push(key);
      }
    }
  }

  saveCache(cache, opts);
  return { cleared: cleared.length, tools: cleared };
}

/**
 * Read-only baseline listing (for the reset CLI's --list mode).
 * @param {object} [opts]
 * @returns {Array<{ tool: string, baseline: string, seenAt: number, lastSeen: number, history: number }>}
 */
export function listBaselines(opts = {}) {
  const cache = loadCache(opts);
  const out = [];
  for (const [tool, entry] of Object.entries(cache)) {
    if (entry && entry.baseline) {
      out.push({
        tool,
        baseline: entry.baseline.description,
        seenAt: entry.baseline.seenAt,
        lastSeen: entry.lastSeen,
        history: Array.isArray(entry.history) ? entry.history.length : 0,
      });
    }
  }
  return out;
}

// ---------------------------------------------------------------------------
// Exported constants (for testing)
// ---------------------------------------------------------------------------
export {
  TTL_MS,
  DRIFT_THRESHOLD,
  CUMULATIVE_DRIFT_THRESHOLD_DEFAULT,
  HISTORY_CAP,
  CACHE_DIR,
  CACHE_FILE,
};