#!/usr/bin/env node
// mcp-live-inspect.mjs — MCP Runtime Inspection Scanner
// Connects to running MCP stdio servers via JSON-RPC 2.0,
// fetches live tool/prompt/resource lists, scans for injection,
// shadowing, and drift. Standalone — not part of scan-orchestrator.
// Zero external dependencies.

import { spawn } from 'node:child_process';
import { createInterface } from 'node:readline';
import { resolve, join } from 'node:path';
import { existsSync, readFileSync } from 'node:fs';
import { homedir } from 'node:os';
import { pathToFileURL } from 'node:url';

import { finding, scannerResult, resetCounter } from './lib/output.mjs';
import { SEVERITY } from './lib/severity.mjs';
import { scanForInjection } from './lib/injection-patterns.mjs';

// ---------------------------------------------------------------------------
// Section 1: MCP Config Discovery
// ---------------------------------------------------------------------------

/**
 * Read a JSON file, returning null on any error.
 * @param {string} filePath
 * @returns {object|null}
 */
function readJsonSafe(filePath) {
  try {
    if (!existsSync(filePath)) return null;
    return JSON.parse(readFileSync(filePath, 'utf8'));
  } catch {
    return null;
  }
}

/**
 * Resolve env variable references in a string (${HOME}, $HOME, ${VAR}).
 * @param {string} str
 * @returns {string}
 */
function resolveEnvVars(str) {
  if (typeof str !== 'string') return str;
  return str
    .replace(/^~(?=[/\\]|$)/, homedir())
    .replace(/\$\{(\w+)\}/g, (_, name) => process.env[name] || '')
    .replace(/\$(\w+)/g, (_, name) => process.env[name] || '');
}

/**
 * Extract MCP server descriptors from a config object.
 * Handles both top-level { mcpServers: {...} } and direct { serverName: {...} } formats.
 * @param {object} config - Parsed JSON config
 * @param {string} sourceFile - Path to the config file
 * @returns {{ servers: object[], skippedSse: number }}
 */
function extractServers(config, sourceFile) {
  const servers = [];
  let skippedSse = 0;

  const block = config?.mcpServers || config;
  if (!block || typeof block !== 'object') return { servers, skippedSse };

  for (const [name, entry] of Object.entries(block)) {
    if (!entry || typeof entry !== 'object') continue;
    // Skip non-server keys (e.g. "enabledPlugins", "permissions")
    if (!entry.command && !entry.url) continue;

    if (entry.url || entry.type === 'sse') {
      skippedSse++;
      continue;
    }

    // Resolve env var references in command, args, env values, and cwd
    const resolvedArgs = (entry.args || []).map(resolveEnvVars);
    const resolvedEnv = {};
    for (const [k, v] of Object.entries(entry.env || {})) {
      resolvedEnv[k] = resolveEnvVars(v);
    }

    servers.push({
      name,
      transport: 'stdio',
      command: resolveEnvVars(entry.command),
      args: resolvedArgs,
      env: resolvedEnv,
      cwd: entry.cwd ? resolveEnvVars(entry.cwd) : null,
      sourceFile,
    });
  }

  return { servers, skippedSse };
}

/**
 * Discover all MCP servers from 6 config locations.
 * @param {string} targetPath - Project root
 * @param {boolean} skipGlobal - Skip ~/.claude/ locations
 * @returns {{ servers: object[], skippedSse: number, configsRead: string[] }}
 */
export function discoverMcpServers(targetPath, skipGlobal = false) {
  const locations = [
    join(targetPath, '.mcp.json'),
    join(targetPath, '.claude', 'settings.json'),
    join(targetPath, 'claude_desktop_config.json'),
  ];

  if (!skipGlobal) {
    const home = homedir();
    locations.push(
      join(home, '.claude', 'settings.json'),
      join(home, '.claude', 'mcp.json'),
      join(home, '.config', 'claude', 'mcp.json'),
    );
  }

  const allServers = [];
  let totalSkippedSse = 0;
  const configsRead = [];
  const seenNames = new Set();

  for (const loc of locations) {
    const config = readJsonSafe(loc);
    if (!config) continue;
    configsRead.push(loc);

    const { servers, skippedSse } = extractServers(config, loc);
    totalSkippedSse += skippedSse;

    for (const s of servers) {
      if (seenNames.has(s.name)) continue; // first wins dedup
      seenNames.add(s.name);
      allServers.push(s);
    }
  }

  return { servers: allServers, skippedSse: totalSkippedSse, configsRead };
}

// ---------------------------------------------------------------------------
// Section 2: JSON-RPC 2.0 Session & Server Inspection
// ---------------------------------------------------------------------------

const DEFAULT_TIMEOUT_MS = 10_000;
const PER_CALL_TIMEOUT_MS = 5_000;
const KILL_GRACE_MS = 500;

/**
 * Create a JSON-RPC 2.0 session over a child process's stdin/stdout.
 * @param {import('child_process').ChildProcess} proc
 * @returns {{ send: function, close: function }}
 */
function createRpcSession(proc) {
  const pending = new Map();
  let nextId = 1;

  const rl = createInterface({ input: proc.stdout });

  rl.on('line', (line) => {
    if (!line.trim()) return;
    let msg;
    try { msg = JSON.parse(line); } catch { return; }
    if (msg.id != null && pending.has(msg.id)) {
      const { resolve: res, reject: rej } = pending.get(msg.id);
      pending.delete(msg.id);
      if (msg.error) {
        const err = new Error(`RPC ${msg.error.code}: ${msg.error.message || 'unknown'}`);
        err.code = msg.error.code;
        rej(err);
      } else {
        res(msg.result);
      }
    }
  });

  proc.stdout.on('close', () => {
    for (const { reject: rej } of pending.values()) {
      rej(new Error('stdout closed'));
    }
    pending.clear();
  });

  /**
   * Send a JSON-RPC message.
   * @param {string} method
   * @param {object} params
   * @param {boolean} expectResponse - false for notifications
   * @returns {Promise<object>}
   */
  function send(method, params = {}, expectResponse = true) {
    const id = expectResponse ? nextId++ : undefined;
    const msg = { jsonrpc: '2.0', method, params };
    if (id !== undefined) msg.id = id;

    try {
      proc.stdin.write(JSON.stringify(msg) + '\n');
    } catch {
      return Promise.reject(new Error('stdin write failed'));
    }

    if (!expectResponse) return Promise.resolve();

    return new Promise((res, rej) => {
      pending.set(id, { resolve: res, reject: rej });
    });
  }

  function close() {
    rl.close();
    pending.clear();
  }

  return { send, close };
}

/**
 * Race a promise against a timeout.
 * @param {Promise} promise
 * @param {number} ms
 * @param {string} label
 * @returns {Promise}
 */
function withTimeout(promise, ms, label = 'operation') {
  return new Promise((res, rej) => {
    const timer = setTimeout(() => rej(new Error(`${label} timed out after ${ms}ms`)), ms);
    promise.then(
      (val) => { clearTimeout(timer); res(val); },
      (err) => { clearTimeout(timer); rej(err); },
    );
  });
}

/**
 * Safely send an RPC call, treating MethodNotFound as empty result.
 * @param {object} session - RPC session
 * @param {string} method
 * @returns {Promise<object>}
 */
async function safeSend(session, method) {
  try {
    return await withTimeout(session.send(method, {}), PER_CALL_TIMEOUT_MS, method);
  } catch (err) {
    if (err.code === -32601) return {}; // MethodNotFound → empty
    throw err;
  }
}

/**
 * Kill a child process with grace period.
 * @param {import('child_process').ChildProcess} proc
 */
function killProcess(proc) {
  try { proc.kill('SIGTERM'); } catch { /* already dead */ }
  setTimeout(() => {
    try { proc.kill('SIGKILL'); } catch { /* already dead */ }
  }, KILL_GRACE_MS);
}

/**
 * Spawn an MCP server, initialize it, fetch tools/prompts/resources.
 * @param {object} descriptor - Server descriptor from discovery
 * @param {number} timeoutMs - Global timeout per server
 * @returns {Promise<object|null>} Inspection result or null on failure
 */
export async function inspectServer(descriptor, timeoutMs = DEFAULT_TIMEOUT_MS) {
  const start = Date.now();

  return new Promise((outerResolve) => {
    let proc;
    let session;
    let globalTimer;
    let resolved = false;

    function finish(result) {
      if (resolved) return;
      resolved = true;
      clearTimeout(globalTimer);
      if (session) session.close();
      if (proc) killProcess(proc);
      outerResolve(result);
    }

    // Global timeout
    globalTimer = setTimeout(() => {
      finish({ name: descriptor.name, error: 'timeout', durationMs: Date.now() - start });
    }, timeoutMs);

    try {
      const mergedEnv = { ...process.env, ...descriptor.env };
      const spawnOpts = {
        stdio: ['pipe', 'pipe', 'pipe'],
        env: mergedEnv,
      };
      if (descriptor.cwd) spawnOpts.cwd = descriptor.cwd;

      proc = spawn(descriptor.command, descriptor.args, spawnOpts);
    } catch (err) {
      finish({ name: descriptor.name, error: `spawn failed: ${err.message}`, durationMs: Date.now() - start });
      return;
    }

    // Capture stderr for diagnostics
    let stderrBuf = '';
    proc.stderr.on('data', (chunk) => {
      if (stderrBuf.length < 500) stderrBuf += chunk.toString();
    });

    proc.on('error', (err) => {
      finish({ name: descriptor.name, error: `process error: ${err.message}`, durationMs: Date.now() - start });
    });

    session = createRpcSession(proc);

    // Run the inspection protocol
    (async () => {
      // 1. Initialize
      const initResult = await withTimeout(
        session.send('initialize', {
          protocolVersion: '2024-11-05',
          capabilities: {},
          clientInfo: { name: 'mcp-live-inspect', version: '1.0.0' },
        }),
        PER_CALL_TIMEOUT_MS,
        'initialize',
      );

      // 2. Send initialized notification
      await session.send('notifications/initialized', {}, false);

      // 3. Fetch lists concurrently
      const [toolsResult, promptsResult, resourcesResult] = await Promise.allSettled([
        safeSend(session, 'tools/list'),
        safeSend(session, 'prompts/list'),
        safeSend(session, 'resources/list'),
      ]);

      finish({
        name: descriptor.name,
        serverInfo: initResult?.serverInfo || null,
        protocolVersion: initResult?.protocolVersion || null,
        tools: toolsResult.status === 'fulfilled' ? (toolsResult.value?.tools || []) : [],
        prompts: promptsResult.status === 'fulfilled' ? (promptsResult.value?.prompts || []) : [],
        resources: resourcesResult.status === 'fulfilled' ? (resourcesResult.value?.resources || []) : [],
        toolsError: toolsResult.status === 'rejected' ? toolsResult.reason.message : null,
        stderr: stderrBuf.trim().slice(0, 200) || null,
        durationMs: Date.now() - start,
        error: null,
      });
    })().catch((err) => {
      finish({
        name: descriptor.name,
        error: `protocol error: ${err.message}`,
        stderr: stderrBuf.trim().slice(0, 200) || null,
        durationMs: Date.now() - start,
      });
    });
  });
}

// ---------------------------------------------------------------------------
// Section 3: Tool Description Injection Analysis
// ---------------------------------------------------------------------------

const URL_IN_DESC_RE = /https?:\/\/\S+|\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b/;

/**
 * Analyze tool descriptions from a server for injection patterns.
 * @param {string} serverName
 * @param {object[]} tools
 * @param {object[]} findings - Mutated: findings pushed here
 */
function analyzeToolDescriptions(serverName, tools, findings) {
  for (const tool of tools) {
    const desc = tool.description || '';
    const name = tool.name || '';

    // Injection pattern scan on description
    if (desc) {
      const result = scanForInjection(desc);
      if (result.found) {
        findings.push(finding({
          scanner: 'MCI',
          severity: result.severity,
          title: `Tool description injection: ${serverName}.${name}`,
          description: `Live tool description contains injection patterns: ${result.patterns.map(p => p.label).join(', ')}`,
          evidence: `[${serverName}] ${name}: ${desc.slice(0, 200)}${desc.length > 200 ? '...' : ''}`,
          owasp: 'MCP03, MCP06',
          recommendation: `Review and sanitize tool description for "${name}" in MCP server "${serverName}". Remove any LLM-directed instructions from tool descriptions.`,
        }));
      }

      // Excessive length
      if (desc.length > 500) {
        findings.push(finding({
          scanner: 'MCI',
          severity: SEVERITY.MEDIUM,
          title: `Excessive tool description length: ${serverName}.${name}`,
          description: `Tool description is ${desc.length} characters (threshold: 500). Long descriptions may contain hidden instructions.`,
          evidence: `[${serverName}] ${name}: ${desc.length} chars`,
          owasp: 'MCP03',
          recommendation: `Review tool description for "${name}" — descriptions over 500 chars are suspicious for embedded instructions.`,
        }));
      }

      // URL/IP in description
      if (URL_IN_DESC_RE.test(desc)) {
        findings.push(finding({
          scanner: 'MCI',
          severity: SEVERITY.HIGH,
          title: `URL/IP in tool description: ${serverName}.${name}`,
          description: `Tool description contains a URL or IP address, which may indicate data exfiltration or tool poisoning.`,
          evidence: `[${serverName}] ${name}: ${desc.slice(0, 200)}`,
          owasp: 'MCP03',
          recommendation: `Investigate why tool "${name}" references external URLs in its description. This is a tool poisoning signal.`,
        }));
      }
    }

    // Injection scan on tool name itself
    if (name) {
      const nameResult = scanForInjection(name);
      if (nameResult.found) {
        findings.push(finding({
          scanner: 'MCI',
          severity: SEVERITY.HIGH,
          title: `Suspicious tool name: ${serverName}.${name}`,
          description: `Tool name contains injection patterns: ${nameResult.patterns.map(p => p.label).join(', ')}`,
          evidence: `[${serverName}] name: ${name}`,
          owasp: 'MCP03, MCP06',
          recommendation: `Tool name "${name}" in server "${serverName}" contains suspicious patterns. Investigate the server's tool registration.`,
        }));
      }
    }
  }
}

// ---------------------------------------------------------------------------
// Section 4: Tool Shadowing Detection
// ---------------------------------------------------------------------------

/**
 * Detect tool shadowing — same tool name across multiple servers.
 * @param {object[]} serverResults - Array of successful inspection results
 * @param {object[]} findings - Mutated
 */
function detectToolShadowing(serverResults, findings) {
  const toolMap = new Map(); // toolName → [{ serverName, description }]

  for (const sr of serverResults) {
    if (!sr.tools) continue;
    for (const tool of sr.tools) {
      const name = tool.name || '';
      if (!name) continue;
      if (!toolMap.has(name)) toolMap.set(name, []);
      toolMap.get(name).push({
        serverName: sr.name,
        description: (tool.description || '').slice(0, 100),
      });
    }
  }

  for (const [toolName, entries] of toolMap) {
    if (entries.length < 2) continue;

    const descriptions = entries.map(e => e.description);
    const allSame = descriptions.every(d => d === descriptions[0]);
    const serverNames = entries.map(e => e.serverName).join(', ');

    findings.push(finding({
      scanner: 'MCI',
      severity: allSame ? SEVERITY.MEDIUM : SEVERITY.HIGH,
      title: `Tool shadowing: "${toolName}" in ${entries.length} servers`,
      description: allSame
        ? `Tool "${toolName}" is defined in ${entries.length} servers with identical descriptions. Likely redundant config.`
        : `Tool "${toolName}" is defined in ${entries.length} servers with DIFFERENT descriptions. One may be impersonating the other.`,
      evidence: entries.map(e => `${e.serverName}: "${e.description}"`).join(' | '),
      owasp: 'MCP09',
      recommendation: allSame
        ? `Review whether "${toolName}" should be served by multiple servers (${serverNames}). Consider consolidating.`
        : `PRIORITY: Different descriptions for "${toolName}" across servers (${serverNames}). Determine which is authoritative and remove the impersonator.`,
    }));
  }
}

// ---------------------------------------------------------------------------
// Section 5: Description Drift Detection (minimal v2.8.0)
// ---------------------------------------------------------------------------

// TODO: S5 — integrate with diff-engine for cross-run description drift detection.
// For v2.8.0, drift detection is limited to tool count comparison when config
// declares expectedToolCount (uncommon). Full drift requires baseline storage.

// ---------------------------------------------------------------------------
// Section 6: Public API + CLI
// ---------------------------------------------------------------------------

/**
 * Connect to all MCP stdio servers discovered from config locations
 * and scan live tool descriptions for security issues.
 *
 * @param {string} targetPath - Project root to resolve relative config paths
 * @param {object} [options]
 * @param {number} [options.timeoutMs=10000] - Per-server timeout
 * @param {boolean} [options.skipGlobal=false] - Skip ~/.claude/ config locations
 * @returns {Promise<object>} scannerResult envelope with meta
 */
export async function inspect(targetPath, options = {}) {
  const start = Date.now();
  const timeoutMs = options.timeoutMs || DEFAULT_TIMEOUT_MS;
  const skipGlobal = options.skipGlobal || false;

  resetCounter();

  // Discover servers
  const discovery = discoverMcpServers(targetPath, skipGlobal);
  const { servers, skippedSse, configsRead } = discovery;

  if (servers.length === 0) {
    const result = scannerResult('mcp-live-inspect', 'ok', [], 0, Date.now() - start);
    result.meta = {
      servers_discovered: 0,
      servers_contacted: 0,
      servers_skipped_sse: skippedSse,
      servers_timed_out: 0,
      servers_failed: 0,
      configs_read: configsRead,
      server_details: [],
    };
    return result;
  }

  // Inspect each server sequentially (avoid spawning many processes at once)
  const allFindings = [];
  const successfulResults = [];
  let contacted = 0;
  let timedOut = 0;
  let failed = 0;
  const serverDetails = [];

  for (const descriptor of servers) {
    const inspResult = await inspectServer(descriptor, timeoutMs);

    if (inspResult.error === 'timeout') {
      timedOut++;
      serverDetails.push({ name: descriptor.name, status: 'timeout', tools: 0 });
      allFindings.push(finding({
        scanner: 'MCI',
        severity: SEVERITY.INFO,
        title: `Server timeout: ${descriptor.name}`,
        description: `MCP server "${descriptor.name}" did not respond within ${timeoutMs}ms. Command: ${descriptor.command} ${(descriptor.args || []).join(' ')}`,
        owasp: 'MCP08',
        recommendation: `Verify that "${descriptor.name}" can start independently. Check command, args, and required env vars.`,
      }));
    } else if (inspResult.error) {
      failed++;
      serverDetails.push({ name: descriptor.name, status: 'failed', error: inspResult.error, tools: 0 });
    } else {
      contacted++;
      const toolCount = inspResult.tools?.length || 0;
      serverDetails.push({
        name: descriptor.name,
        status: 'ok',
        tools: toolCount,
        prompts: inspResult.prompts?.length || 0,
        resources: inspResult.resources?.length || 0,
        durationMs: inspResult.durationMs,
      });
      successfulResults.push(inspResult);

      // Analyze tool descriptions
      if (inspResult.tools && inspResult.tools.length > 0) {
        analyzeToolDescriptions(descriptor.name, inspResult.tools, allFindings);
      }
    }
  }

  // Cross-server analysis: tool shadowing
  if (successfulResults.length >= 2) {
    detectToolShadowing(successfulResults, allFindings);
  }

  const durationMs = Date.now() - start;
  const result = scannerResult('mcp-live-inspect', 'ok', allFindings, contacted, durationMs);
  result.meta = {
    servers_discovered: servers.length,
    servers_contacted: contacted,
    servers_skipped_sse: skippedSse,
    servers_timed_out: timedOut,
    servers_failed: failed,
    configs_read: configsRead,
    server_details: serverDetails,
  };

  return result;
}

// ---------------------------------------------------------------------------
// CLI entry point
// ---------------------------------------------------------------------------

function parseCliArgs(argv) {
  const args = { target: null, timeoutMs: DEFAULT_TIMEOUT_MS, skipGlobal: false };
  for (let i = 2; i < argv.length; i++) {
    if (argv[i] === '--timeout' && argv[i + 1]) {
      args.timeoutMs = parseInt(argv[++i], 10) || DEFAULT_TIMEOUT_MS;
    } else if (argv[i] === '--skip-global') {
      args.skipGlobal = true;
    } else if (!args.target) {
      args.target = argv[i];
    }
  }
  return args;
}

if (import.meta.url === pathToFileURL(resolve(process.argv[1] || '')).href) {
  const args = parseCliArgs(process.argv);
  const target = resolve(args.target || '.');

  inspect(target, { timeoutMs: args.timeoutMs, skipGlobal: args.skipGlobal })
    .then((result) => {
      process.stdout.write(JSON.stringify(result, null, 2) + '\n');
      const v = result.counts || {};
      if ((v.critical || 0) >= 1) process.exit(2);
      if ((v.high || 0) >= 1) process.exit(1);
      process.exit(0);
    })
    .catch((err) => {
      process.stderr.write(`Fatal: ${err.message}\n`);
      process.exit(1);
    });
}
