contextguard 0.1.3

package/src/mcp-security-wrapper.ts

@@ -0,0 +1,527 @@
+#!/usr/bin/env node
+
+/**
+ * Copyright (c) 2025 Amir Mironi
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { spawn, ChildProcess } from "child_process";
+import * as fs from "fs";
+import { createHash } from "crypto";
+
+interface SecurityConfig {
+  maxToolCallsPerMinute?: number;
+  blockedPatterns?: string[];
+  allowedFilePaths?: string[];
+  alertThreshold?: number;
+  enablePromptInjectionDetection?: boolean;
+  enableSensitiveDataDetection?: boolean;
+}
+
+interface SecurityEvent {
+  timestamp: string;
+  eventType: string;
+  severity: "LOW" | "MEDIUM" | "HIGH" | "CRITICAL";
+  details: Record<string, unknown>;
+  sessionId: string;
+}
+
+interface MCPMessage {
+  jsonrpc: string;
+  id?: string | number;
+  method?: string;
+  params?: {
+    name?: string;
+    arguments?: Record<string, string>;
+    path?: string;
+    filePath?: string;
+    [key: string]: unknown;
+  };
+  result?: unknown;
+  error?: unknown;
+}
+
+class SecurityPolicy {
+  private config: SecurityConfig;
+  private sensitiveDataPatterns: RegExp[];
+  private promptInjectionPatterns: RegExp[];
+
+  constructor(config: SecurityConfig) {
+    this.config = {
+      maxToolCallsPerMinute: 30,
+      blockedPatterns: [],
+      allowedFilePaths: [],
+      alertThreshold: 5,
+      enablePromptInjectionDetection: true,
+      enableSensitiveDataDetection: true,
+      ...config,
+    };
+
+    // Sensitive data patterns
+    this.sensitiveDataPatterns = [
+      /(?:password|secret|api[_-]?key|token)\s*[:=]\s*['"]?[\w\-.]+['"]?/gi,
+      /\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b/g, // Email
+      /\b\d{3}-\d{2}-\d{4}\b/g, // SSN
+      /sk-[a-zA-Z0-9]{48}/g, // OpenAI API keys
+      /ghp_[a-zA-Z0-9]{36}/g, // GitHub tokens
+      /AKIA[0-9A-Z]{16}/g, // AWS Access Keys
+    ];
+
+    // Prompt injection patterns
+    this.promptInjectionPatterns = [
+      /ignore\s+(previous|all)\s+(instructions|prompts)/gi,
+      /system:\s*you\s+are\s+now/gi,
+      /forget\s+(everything|all)/gi,
+      /new\s+instructions:/gi,
+      /\[INST\].*?\[\/INST\]/gs,
+      /<\|im_start\|>/g,
+      /disregard\s+previous/gi,
+      /override\s+previous/gi,
+    ];
+  }
+
+  checkPromptInjection(text: string): string[] {
+    if (!this.config.enablePromptInjectionDetection) return [];
+    const violations: string[] = [];
+    for (const pattern of this.promptInjectionPatterns) {
+      const matches = text.match(pattern);
+      if (matches) {
+        violations.push(
+          `Potential prompt injection detected: "${matches[0].substring(
+            0,
+            50
+          )}..."`
+        );
+      }
+    }
+    return violations;
+  }
+
+  checkSensitiveData(text: string): string[] {
+    if (!this.config.enableSensitiveDataDetection) return [];
+    const violations: string[] = [];
+    for (const pattern of this.sensitiveDataPatterns) {
+      const matches = text.match(pattern);
+      if (matches) {
+        violations.push(
+          `Sensitive data pattern detected (redacted): ${pattern.source.substring(
+            0,
+            30
+          )}...`
+        );
+      }
+    }
+    return violations;
+  }
+
+  checkFileAccess(filePath: string): string[] {
+    const violations: string[] = [];
+    if (filePath.includes("..")) {
+      violations.push(`Path traversal attempt detected: ${filePath}`);
+    }
+    const dangerousPaths = [
+      "/etc",
+      "/root",
+      "/sys",
+      "/proc",
+      "C:\\Windows\\System32",
+    ];
+    if (dangerousPaths.some((dangerous) => filePath.startsWith(dangerous))) {
+      violations.push(`Access to dangerous path detected: ${filePath}`);
+    }
+    if (
+      this.config.allowedFilePaths &&
+      this.config.allowedFilePaths.length > 0
+    ) {
+      const isAllowed = this.config.allowedFilePaths.some((allowed) =>
+        filePath.startsWith(allowed)
+      );
+      if (!isAllowed) {
+        violations.push(`File path not in allowed list: ${filePath}`);
+      }
+    }
+    return violations;
+  }
+
+  checkRateLimit(timestamps: number[]): boolean {
+    const oneMinuteAgo = Date.now() - 60000;
+    const recentCalls = timestamps.filter((t) => t > oneMinuteAgo);
+    return recentCalls.length < (this.config.maxToolCallsPerMinute || 30);
+  }
+}
+
+class SecurityLogger {
+  private logFile: string;
+  private events: SecurityEvent[] = [];
+
+  constructor(logFile: string = "mcp_security.log") {
+    this.logFile = logFile;
+  }
+
+  logEvent(
+    eventType: string,
+    severity: SecurityEvent["severity"],
+    details: Record<string, unknown>,
+    sessionId: string
+  ): void {
+    const event: SecurityEvent = {
+      timestamp: new Date().toISOString(),
+      eventType,
+      severity,
+      details,
+      sessionId,
+    };
+    this.events.push(event);
+    fs.appendFileSync(this.logFile, JSON.stringify(event) + "\n");
+    if (severity === "HIGH" || severity === "CRITICAL") {
+      console.error(
+        `[SECURITY ALERT] ${eventType}: ${JSON.stringify(details)}`
+      );
+    }
+  }
+
+  getStatistics(): Record<string, unknown> {
+    return {
+      totalEvents: this.events.length,
+      eventsByType: this.countByField("eventType"),
+      eventsBySeverity: this.countByField("severity"),
+      recentEvents: this.events.slice(-10),
+    };
+  }
+
+  private countByField(field: keyof SecurityEvent): Record<string, number> {
+    const counts: Record<string, number> = {};
+    for (const event of this.events) {
+      const value = String(event[field]);
+      counts[value] = (counts[value] || 0) + 1;
+    }
+    return counts;
+  }
+}
+
+class MCPSecurityWrapper {
+  private serverCommand: string[];
+  private policy: SecurityPolicy;
+  private logger: SecurityLogger;
+  private process: ChildProcess | null = null;
+  private toolCallTimestamps: number[] = [];
+  private sessionId: string;
+  private clientMessageBuffer: string = "";
+  private serverMessageBuffer: string = "";
+
+  constructor(
+    serverCommand: string[],
+    policy: SecurityPolicy,
+    logger: SecurityLogger
+  ) {
+    this.serverCommand = serverCommand;
+    this.policy = policy;
+    this.logger = logger;
+    this.sessionId = createHash("md5")
+      .update(Date.now().toString())
+      .digest("hex")
+      .substring(0, 8);
+  }
+
+  async start(): Promise<void> {
+    this.process = spawn(this.serverCommand[0], this.serverCommand.slice(1), {
+      stdio: ["pipe", "pipe", "pipe"],
+    });
+
+    if (!this.process.stdout || !this.process.stdin || !this.process.stderr) {
+      throw new Error("Failed to create child process streams");
+    }
+
+    this.logger.logEvent(
+      "SERVER_START",
+      "LOW",
+      {
+        command: this.serverCommand.join(" "),
+        pid: this.process.pid,
+      },
+      this.sessionId
+    );
+
+    this.process.stderr.pipe(process.stderr);
+
+    this.process.stdout.on("data", (data: Buffer) => {
+      const output = data.toString();
+      this.handleServerOutput(output);
+    });
+
+    process.stdin.on("data", (data: Buffer) => {
+      const input = data.toString();
+      this.handleClientInput(input);
+    });
+
+    this.process.on("exit", (code) => {
+      this.logger.logEvent(
+        "SERVER_EXIT",
+        "MEDIUM",
+        {
+          exitCode: code,
+        },
+        this.sessionId
+      );
+      console.error("\n=== MCP Security Statistics ===");
+      console.error(JSON.stringify(this.logger.getStatistics(), null, 2));
+      // Use setImmediate to allow pending I/O to complete
+      setImmediate(() => {
+        process.exit(code || 0);
+      });
+    });
+
+    this.process.on("error", (err) => {
+      this.logger.logEvent(
+        "SERVER_ERROR",
+        "HIGH",
+        {
+          error: err.message,
+        },
+        this.sessionId
+      );
+      console.error("Server process error:", err);
+    });
+  }
+
+  private handleClientInput(input: string): void {
+    this.clientMessageBuffer += input;
+    const lines = this.clientMessageBuffer.split("\n");
+    this.clientMessageBuffer = lines.pop() || "";
+    for (const line of lines) {
+      if (line.trim()) {
+        this.processClientMessage(line);
+      }
+    }
+  }
+
+  private processClientMessage(line: string): void {
+    try {
+      const message: MCPMessage = JSON.parse(line);
+      this.logger.logEvent(
+        "CLIENT_REQUEST",
+        "LOW",
+        {
+          method: message.method,
+          id: message.id,
+        },
+        this.sessionId
+      );
+
+      const violations: string[] = [];
+      let shouldBlock = false;
+
+      if (message.method === "tools/call") {
+        const now = Date.now();
+        this.toolCallTimestamps.push(now);
+
+        // Clean up old timestamps to prevent memory leak
+        const oneMinuteAgo = now - 60000;
+        this.toolCallTimestamps = this.toolCallTimestamps.filter(
+          (t) => t > oneMinuteAgo
+        );
+
+        if (!this.policy.checkRateLimit(this.toolCallTimestamps)) {
+          violations.push("Rate limit exceeded for tool calls");
+          shouldBlock = true;
+          this.logger.logEvent(
+            "RATE_LIMIT_EXCEEDED",
+            "HIGH",
+            {
+              method: message.method,
+              toolName: message.params?.name,
+            },
+            this.sessionId
+          );
+        }
+
+        const paramsStr = JSON.stringify(message.params);
+        const injectionViolations = this.policy.checkPromptInjection(paramsStr);
+        violations.push(...injectionViolations);
+        const sensitiveViolations = this.policy.checkSensitiveData(paramsStr);
+        violations.push(...sensitiveViolations);
+
+        // Check multiple possible file path parameter locations
+        const filePathParams = [
+          message.params?.arguments?.path,
+          message.params?.arguments?.filePath,
+          message.params?.arguments?.file,
+          message.params?.arguments?.directory,
+          message.params?.path,
+          message.params?.filePath,
+        ].filter((path): path is string => typeof path === 'string');
+
+        for (const filePath of filePathParams) {
+          const fileViolations = this.policy.checkFileAccess(filePath);
+          violations.push(...fileViolations);
+        }
+
+        this.logger.logEvent(
+          "TOOL_CALL",
+          violations.length > 0 ? "HIGH" : "LOW",
+          {
+            toolName: message.params?.name,
+            hasViolations: violations.length > 0,
+            violations,
+          },
+          this.sessionId
+        );
+      }
+
+      if (violations.length > 0) {
+        this.logger.logEvent(
+          "SECURITY_VIOLATION",
+          "CRITICAL",
+          {
+            violations,
+            message: message,
+            blocked: shouldBlock,
+          },
+          this.sessionId
+        );
+        console.error(
+          `\n⚠️  SECURITY VIOLATIONS DETECTED:\n${violations.join("\n")}\n`
+        );
+
+        if (shouldBlock) {
+          console.error("🚫 REQUEST BLOCKED\n");
+          // Send error response back to client
+          if (message.id !== undefined) {
+            const errorResponse: MCPMessage = {
+              jsonrpc: message.jsonrpc,
+              id: message.id,
+              error: {
+                code: -32000,
+                message: "Security violation: Request blocked",
+                data: { violations },
+              },
+            };
+            process.stdout.write(JSON.stringify(errorResponse) + "\n");
+          }
+          return; // Don't forward to server
+        }
+      }
+
+      if (this.process && this.process.stdin) {
+        this.process.stdin.write(line + "\n");
+      }
+    } catch (err) {
+      this.logger.logEvent(
+        "PARSE_ERROR",
+        "MEDIUM",
+        {
+          error: err instanceof Error ? err.message : String(err),
+          line: line.substring(0, 100),
+        },
+        this.sessionId
+      );
+      console.error(`Failed to parse client message: ${err}`);
+      // Forward unparseable messages
+      if (this.process && this.process.stdin) {
+        this.process.stdin.write(line + "\n");
+      }
+    }
+  }
+
+  private handleServerOutput(output: string): void {
+    // Forward output immediately
+    process.stdout.write(output);
+
+    // Buffer and parse for logging
+    this.serverMessageBuffer += output;
+    const lines = this.serverMessageBuffer.split("\n");
+    this.serverMessageBuffer = lines.pop() || "";
+
+    for (const line of lines) {
+      if (line.trim()) {
+        try {
+          const message: MCPMessage = JSON.parse(line);
+          this.logger.logEvent(
+            "SERVER_RESPONSE",
+            "LOW",
+            {
+              id: message.id,
+              hasError: !!message.error,
+            },
+            this.sessionId
+          );
+        } catch (err) {
+          // Log parse errors for server output
+          this.logger.logEvent(
+            "SERVER_PARSE_ERROR",
+            "LOW",
+            {
+              error: err instanceof Error ? err.message : String(err),
+              line: line.substring(0, 100),
+            },
+            this.sessionId
+          );
+        }
+      }
+    }
+  }
+}
+
+async function main() {
+  const args = process.argv.slice(2);
+  if (args.length === 0 || args.includes("--help")) {
+    console.log(`
+MCP Security Wrapper - MVP
+
+Usage:
+  npx ts-node mcp-security-wrapper.ts --server "node server.js" [--config security.json]
+
+Options:
+  --server <command>    Command to start the MCP server (required)
+  --config <file>       Path to security config JSON file (optional)
+  --help               Show this help message
+
+Example:
+  npx ts-node mcp-security-wrapper.ts --server "node server.js" --config security.json
+    `);
+    process.exit(0);
+  }
+
+  let serverCommand: string = "";
+  let configFile: string = "";
+
+  for (let i = 0; i < args.length; i++) {
+    if (args[i] === "--server" && args[i + 1]) {
+      serverCommand = args[i + 1];
+      i++;
+    } else if (args[i] === "--config" && args[i + 1]) {
+      configFile = args[i + 1];
+      i++;
+    }
+  }
+
+  if (!serverCommand) {
+    console.error("Error: --server argument is required");
+    process.exit(1);
+  }
+
+  let config: SecurityConfig = {};
+  if (configFile && fs.existsSync(configFile)) {
+    config = JSON.parse(fs.readFileSync(configFile, "utf-8"));
+  }
+
+  const policy = new SecurityPolicy(config);
+  const logger = new SecurityLogger();
+  const wrapper = new MCPSecurityWrapper(
+    serverCommand.split(" "),
+    policy,
+    logger
+  );
+  await wrapper.start();
+}
+
+if (require.main === module) {
+  main().catch((err) => {
+    console.error("Fatal error:", err);
+    process.exit(1);
+  });
+}
+
+export { MCPSecurityWrapper, SecurityPolicy, SecurityLogger };