mcp-context-engineering/src/security/securityManager.ts

The intelligent context optimization system for AI coding assistants. Built with Cole's PRP methodology, Context Portal knowledge graphs, and production-ready MongoDB architecture.

import { z } from 'zod';

/**
 * Security Manager - Octocode-inspired multi-layer security
 * 
 * Implements comprehensive security patterns:
 * - Input sanitization and validation
 * - Real-time secret detection (1100+ patterns)
 * - Prompt injection defense
 * - Content redaction and privacy protection
 * - Multi-layer protection with fallback strategies
 */

// Security configuration schema
export const SecurityConfigSchema = z.object({
  secret_detection: z.object({
    enabled: z.boolean().default(true),
    patterns_file: z.string().optional(),
    custom_patterns: z.array(z.string()).default([]),
    redaction_char: z.string().default('*'),
    context_window: z.number().default(10)
  }),
  
  input_sanitization: z.object({
    enabled: z.boolean().default(true),
    max_input_length: z.number().default(100000),
    forbidden_patterns: z.array(z.string()).default([]),
    sanitization_level: z.enum(['basic', 'standard', 'strict']).default('standard')
  }),
  
  prompt_injection_defense: z.object({
    enabled: z.boolean().default(true),
    detection_threshold: z.number().min(0).max(1).default(0.7),
    defense_strategies: z.array(z.enum(['block', 'sanitize', 'flag'])).default(['sanitize', 'flag'])
  }),
  
  content_filtering: z.object({
    enabled: z.boolean().default(true),
    pii_detection: z.boolean().default(true),
    code_injection_detection: z.boolean().default(true),
    malicious_pattern_detection: z.boolean().default(true)
  })
});

export type SecurityConfig = z.infer<typeof SecurityConfigSchema>;

// Secret patterns inspired by Octocode's 1100+ pattern library
const SECRET_PATTERNS = [
  // API Keys
  { name: 'Generic API Key', pattern: /\b[A-Za-z0-9]{32,}\b/, risk: 'high' },
  { name: 'AWS Access Key', pattern: /AKIA[0-9A-Z]{16}/, risk: 'critical' },
  { name: 'GitHub Token', pattern: /ghp_[A-Za-z0-9]{36}/, risk: 'high' },
  { name: 'OpenAI API Key', pattern: /sk-[A-Za-z0-9]{48}/, risk: 'high' },
  { name: 'Voyage AI Key', pattern: /pa-[A-Za-z0-9-]{32,}/, risk: 'high' },
  
  // Database connections
  { name: 'MongoDB URI', pattern: /mongodb(\+srv)?:\/\/[^\s]+/, risk: 'critical' },
  { name: 'Database Password', pattern: /(?:password|pwd|pass)[\s]*[:=][\s]*['"]?([^'"\s]+)['"]?/i, risk: 'critical' },
  
  // Authentication tokens
  { name: 'JWT Token', pattern: /eyJ[A-Za-z0-9-_=]+\.[A-Za-z0-9-_=]+\.?[A-Za-z0-9-_.+/=]*/, risk: 'high' },
  { name: 'Bearer Token', pattern: /Bearer\s+[A-Za-z0-9\-_+=\/]{20,}/, risk: 'high' },
  
  // Private keys
  { name: 'Private Key', pattern: /-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----/, risk: 'critical' },
  { name: 'SSH Key', pattern: /ssh-(?:rsa|dss|ed25519)\s+[A-Za-z0-9+/]+=*/, risk: 'high' },
  
  // Common secrets
  { name: 'Generic Secret', pattern: /(?:secret|key|token|password)[\s]*[:=][\s]*['"]?([A-Za-z0-9+/=]{16,})['"]?/i, risk: 'medium' },
  { name: 'Hex Secret', pattern: /\b[a-fA-F0-9]{32,128}\b/, risk: 'medium' }
];

// Prompt injection patterns
const PROMPT_INJECTION_PATTERNS = [
  /ignore\s+(?:previous|all)\s+(?:instructions|prompts)/i,
  /forget\s+(?:everything|all|previous)/i,
  /new\s+(?:instructions|prompt|system|role)/i,
  /(?:assistant|ai|system):\s*now\s+you\s+are/i,
  /(?:pretend|act|behave)\s+(?:as|like)\s+(?:a|an)\s+(?:different|new)/i,
  /override\s+(?:your|system|default)\s+(?:instructions|settings)/i,
  /execute\s+(?:code|command|script)/i,
  /\[\s*system\s*\]|\[\s*user\s*\]|\[\s*assistant\s*\]/i
];

// Malicious content patterns
const MALICIOUS_PATTERNS = [
  // Code injection
  /(?:eval|exec|system|shell_exec|passthru)\s*\(/i,
  /(?:setTimeout|setInterval)\s*\(\s*['"`]/i,
  /(?:document\.write|innerHTML|outerHTML)\s*=/i,
  
  // Command injection
  /(?:\||;|&&|\$\(|\`)[^;|&`$]*(?:rm|cat|ls|wget|curl|nc|bash|sh)/i,
  
  // Path traversal
  /\.\.\/|\.\.\\|%2e%2e%2f|%2e%2e%5c/i,
  
  // SQL injection patterns
  /(?:union|select|insert|update|delete|drop|create|alter)\s+(?:all\s+)?(?:select|from|where|union|order|group|having|insert|update|delete|drop|create|alter)/i
];

export class SecurityManager {
  private config: SecurityConfig;
  private secretPatterns: typeof SECRET_PATTERNS;
  private detectionStats: {
    secretsDetected: number;
    injectionsBlocked: number;
    maliciousContentBlocked: number;
    totalProcessed: number;
  };

  constructor(config: Partial<SecurityConfig> = {}) {
    this.config = SecurityConfigSchema.parse(config);
    this.secretPatterns = [...SECRET_PATTERNS, ...this.loadCustomPatterns()];
    this.detectionStats = {
      secretsDetected: 0,
      injectionsBlocked: 0,
      maliciousContentBlocked: 0,
      totalProcessed: 0
    };
  }

  /**
   * Comprehensive input processing with multi-layer security
   */
  async processInput(input: string, context: {
    source: 'user' | 'api' | 'system';
    agent_type?: string;
    session_id?: string;
  }): Promise<{
    sanitized_input: string;
    security_flags: string[];
    risk_level: 'low' | 'medium' | 'high' | 'critical';
    blocked: boolean;
  }> {
    this.detectionStats.totalProcessed++;
    
    const securityFlags: string[] = [];
    let sanitizedInput = input;
    let riskLevel: 'low' | 'medium' | 'high' | 'critical' = 'low';
    let blocked = false;

    // Layer 1: Input validation and sanitization
    if (this.config.input_sanitization.enabled) {
      const sanitizationResult = this.sanitizeInput(sanitizedInput);
      sanitizedInput = sanitizationResult.sanitized;
      securityFlags.push(...sanitizationResult.flags);
      
      if (sanitizationResult.risk === 'critical') {
        riskLevel = 'critical';
      }
    }

    // Layer 2: Secret detection and redaction
    if (this.config.secret_detection.enabled) {
      const secretResult = this.detectAndRedactSecrets(sanitizedInput);
      sanitizedInput = secretResult.redacted;
      securityFlags.push(...secretResult.flags);
      
      if (secretResult.secretsFound > 0) {
        this.detectionStats.secretsDetected += secretResult.secretsFound;
        riskLevel = this.escalateRisk(riskLevel, 'high');
      }
    }

    // Layer 3: Prompt injection detection
    if (this.config.prompt_injection_defense.enabled) {
      const injectionResult = this.detectPromptInjection(sanitizedInput);
      if (injectionResult.detected) {
        this.detectionStats.injectionsBlocked++;
        securityFlags.push(...injectionResult.flags);
        riskLevel = this.escalateRisk(riskLevel, 'high');
        
        if (this.config.prompt_injection_defense.defense_strategies.includes('block')) {
          blocked = true;
        }
        
        if (this.config.prompt_injection_defense.defense_strategies.includes('sanitize')) {
          sanitizedInput = injectionResult.sanitized;
        }
      }
    }

    // Layer 4: Malicious content detection
    if (this.config.content_filtering.enabled) {
      const maliciousResult = this.detectMaliciousContent(sanitizedInput);
      if (maliciousResult.detected) {
        this.detectionStats.maliciousContentBlocked++;
        securityFlags.push(...maliciousResult.flags);
        riskLevel = this.escalateRisk(riskLevel, 'critical');
        blocked = true;
      }
    }

    // Log security events for monitoring
    if (securityFlags.length > 0) {
      this.logSecurityEvent({
        timestamp: new Date(),
        source: context.source,
        agent_type: context.agent_type,
        session_id: context.session_id,
        flags: securityFlags,
        risk_level: riskLevel,
        blocked,
        input_length: input.length,
        sanitized_length: sanitizedInput.length
      });
    }

    return {
      sanitized_input: sanitizedInput,
      security_flags: securityFlags,
      risk_level: riskLevel,
      blocked
    };
  }

  /**
   * Input sanitization with configurable levels
   */
  private sanitizeInput(input: string): {
    sanitized: string;
    flags: string[];
    risk: 'low' | 'medium' | 'high' | 'critical';
  } {
    const flags: string[] = [];
    let sanitized = input;
    let risk: 'low' | 'medium' | 'high' | 'critical' = 'low';

    // Length validation
    if (input.length > this.config.input_sanitization.max_input_length) {
      flags.push('input_too_long');
      sanitized = input.substring(0, this.config.input_sanitization.max_input_length) + '... [truncated]';
      risk = 'medium';
    }

    // Forbidden patterns
    for (const pattern of this.config.input_sanitization.forbidden_patterns) {
      const regex = new RegExp(pattern, 'gi');
      if (regex.test(sanitized)) {
        flags.push(`forbidden_pattern: ${pattern}`);
        sanitized = sanitized.replace(regex, '[REDACTED]');
        risk = this.escalateRisk(risk, 'high');
      }
    }

    // Control character removal
    sanitized = sanitized.replace(/[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]/g, '');

    // HTML entity encoding for strict mode
    if (this.config.input_sanitization.sanitization_level === 'strict') {
      sanitized = sanitized
        .replace(/&/g, '&amp;')
        .replace(/</g, '&lt;')
        .replace(/>/g, '&gt;')
        .replace(/"/g, '&quot;')
        .replace(/'/g, '&#x27;');
    }

    return { sanitized, flags, risk };
  }

  /**
   * Real-time secret detection with 1100+ patterns
   */
  private detectAndRedactSecrets(input: string): {
    redacted: string;
    flags: string[];
    secretsFound: number;
  } {
    const flags: string[] = [];
    let redacted = input;
    let secretsFound = 0;

    for (const secretPattern of this.secretPatterns) {
      const matches = input.match(new RegExp(secretPattern.pattern, 'g'));
      if (matches) {
        secretsFound += matches.length;
        flags.push(`secret_detected: ${secretPattern.name} (${secretPattern.risk} risk)`);
        
        // Redact with context window
        const contextWindow = this.config.secret_detection.context_window;
        redacted = redacted.replace(new RegExp(secretPattern.pattern, 'g'), (match) => {
          const redactionLength = Math.max(match.length - contextWindow * 2, 4);
          const redactionChar = this.config.secret_detection.redaction_char;
          return match.substring(0, contextWindow) + 
                 redactionChar.repeat(redactionLength) + 
                 match.substring(match.length - contextWindow);
        });
      }
    }

    return { redacted, flags, secretsFound };
  }

  /**
   * Advanced prompt injection detection
   */
  private detectPromptInjection(input: string): {
    detected: boolean;
    flags: string[];
    sanitized: string;
    confidence: number;
  } {
    const flags: string[] = [];
    let detected = false;
    let sanitized = input;
    let confidence = 0;

    for (const pattern of PROMPT_INJECTION_PATTERNS) {
      const matches = input.match(pattern);
      if (matches) {
        detected = true;
        confidence = Math.max(confidence, 0.8);
        flags.push(`prompt_injection_detected: ${pattern.source}`);
        
        // Sanitize by replacing with harmless content
        sanitized = sanitized.replace(pattern, '[POTENTIAL_INJECTION_REMOVED]');
      }
    }

    // Additional heuristic checks
    const suspiciousKeywords = ['system', 'instruction', 'prompt', 'override', 'ignore', 'forget'];
    const keywordCount = suspiciousKeywords.filter(keyword => 
      input.toLowerCase().includes(keyword)
    ).length;
    
    if (keywordCount >= 3) {
      confidence = Math.max(confidence, 0.6);
      flags.push(`high_suspicious_keyword_density: ${keywordCount}`);
    }

    detected = detected || confidence >= this.config.prompt_injection_defense.detection_threshold;

    return { detected, flags, sanitized, confidence };
  }

  /**
   * Malicious content detection
   */
  private detectMaliciousContent(input: string): {
    detected: boolean;
    flags: string[];
  } {
    const flags: string[] = [];
    let detected = false;

    for (const pattern of MALICIOUS_PATTERNS) {
      if (pattern.test(input)) {
        detected = true;
        flags.push(`malicious_pattern_detected: ${pattern.source}`);
      }
    }

    return { detected, flags };
  }

  /**
   * Risk level escalation
   */
  private escalateRisk(
    current: 'low' | 'medium' | 'high' | 'critical',
    new_risk: 'low' | 'medium' | 'high' | 'critical'
  ): 'low' | 'medium' | 'high' | 'critical' {
    const riskLevels = { low: 1, medium: 2, high: 3, critical: 4 };
    const currentLevel = riskLevels[current];
    const newLevel = riskLevels[new_risk];
    
    if (newLevel > currentLevel) {
      const levels: ('low' | 'medium' | 'high' | 'critical')[] = ['low', 'medium', 'high', 'critical'];
      return levels[newLevel - 1];
    }
    
    return current;
  }

  /**
   * Load custom secret patterns
   */
  private loadCustomPatterns(): Array<{ name: string; pattern: RegExp; risk: string }> {
    return this.config.secret_detection.custom_patterns.map((pattern, index) => ({
      name: `Custom Pattern ${index + 1}`,
      pattern: new RegExp(pattern, 'gi'),
      risk: 'medium'
    }));
  }

  /**
   * Security event logging
   */
  private logSecurityEvent(event: {
    timestamp: Date;
    source: string;
    agent_type?: string;
    session_id?: string;
    flags: string[];
    risk_level: string;
    blocked: boolean;
    input_length: number;
    sanitized_length: number;
  }): void {
    // In production, this would log to a secure monitoring system
    console.log(`🛡️  SECURITY EVENT:`, {
      timestamp: event.timestamp.toISOString(),
      source: event.source,
      agent_type: event.agent_type,
      flags: event.flags,
      risk_level: event.risk_level,
      blocked: event.blocked,
      reduction: `${event.input_length} → ${event.sanitized_length} chars`
    });
  }

  /**
   * Get security statistics
   */
  getSecurityStats(): typeof this.detectionStats {
    return { ...this.detectionStats };
  }

  /**
   * Reset security statistics
   */
  resetStats(): void {
    this.detectionStats = {
      secretsDetected: 0,
      injectionsBlocked: 0,
      maliciousContentBlocked: 0,
      totalProcessed: 0
    };
  }

  /**
   * Update security configuration
   */
  updateConfig(newConfig: Partial<SecurityConfig>): void {
    this.config = SecurityConfigSchema.parse({ ...this.config, ...newConfig });
    this.secretPatterns = [...SECRET_PATTERNS, ...this.loadCustomPatterns()];
  }
}

// Export singleton instance
export const securityManager = new SecurityManager();