smartproxy/ts/core/utils/security-utils.ts

import * as plugins from '../../plugins.js';
import { 
  matchIpPattern, 
  ipToNumber,
  matchIpCidr
} from './route-utils.js';

/**
 * Security utilities for IP validation, rate limiting, 
 * authentication, and other security features
 */

/**
 * Result of IP validation
 */
export interface IIpValidationResult {
  allowed: boolean;
  reason?: string;
}

/**
 * IP connection tracking information
 */
export interface IIpConnectionInfo {
  connections: Set<string>;  // ConnectionIDs
  timestamps: number[];      // Connection timestamps
  ipVariants: string[];      // Normalized IP variants (e.g., ::ffff:127.0.0.1 and 127.0.0.1)
}

/**
 * Rate limit tracking
 */
export interface IRateLimitInfo {
  count: number;
  expiry: number;
}

/**
 * Logger interface for security utilities
 */
export interface ISecurityLogger {
  info: (message: string, ...args: any[]) => void;
  warn: (message: string, ...args: any[]) => void;
  error: (message: string, ...args: any[]) => void;
  debug?: (message: string, ...args: any[]) => void;
}

/**
 * Normalize IP addresses for comparison
 * Handles IPv4-mapped IPv6 addresses (::ffff:127.0.0.1)
 * 
 * @param ip IP address to normalize
 * @returns Array of equivalent IP representations
 */
export function normalizeIP(ip: string): string[] {
  if (!ip) return [];
  
  // Handle IPv4-mapped IPv6 addresses (::ffff:127.0.0.1)
  if (ip.startsWith('::ffff:')) {
    const ipv4 = ip.slice(7);
    return [ip, ipv4];
  }
  
  // Handle IPv4 addresses by also checking IPv4-mapped form
  if (/^\d{1,3}(\.\d{1,3}){3}$/.test(ip)) {
    return [ip, `::ffff:${ip}`];
  }
  
  return [ip];
}

/**
 * Check if an IP is authorized based on allow and block lists
 *
 * @param ip - The IP address to check
 * @param allowedIPs - Array of allowed IP patterns
 * @param blockedIPs - Array of blocked IP patterns
 * @returns Whether the IP is authorized
 */
export function isIPAuthorized(
  ip: string, 
  allowedIPs: string[] = ['*'], 
  blockedIPs: string[] = []
): boolean {
  // Skip IP validation if no rules
  if (!ip || (allowedIPs.length === 0 && blockedIPs.length === 0)) {
    return true;
  }

  // First check if IP is blocked - blocked IPs take precedence
  if (blockedIPs.length > 0) {
    for (const pattern of blockedIPs) {
      if (matchIpPattern(pattern, ip)) {
        return false;
      }
    }
  }

  // If allowed IPs list has wildcard, all non-blocked IPs are allowed
  if (allowedIPs.includes('*')) {
    return true;
  }

  // Then check if IP is allowed in the explicit allow list
  if (allowedIPs.length > 0) {
    for (const pattern of allowedIPs) {
      if (matchIpPattern(pattern, ip)) {
        return true;
      }
    }
    // If allowedIPs is specified but no match, deny access
    return false;
  }

  // Default allow if no explicit allow list
  return true;
}

/**
 * Check if an IP exceeds maximum connections
 *
 * @param ip - The IP address to check
 * @param ipConnectionsMap - Map of IPs to connection info
 * @param maxConnectionsPerIP - Maximum allowed connections per IP
 * @returns Result with allowed status and reason if blocked
 */
export function checkMaxConnections(
  ip: string,
  ipConnectionsMap: Map<string, IIpConnectionInfo>,
  maxConnectionsPerIP: number
): IIpValidationResult {
  if (!ipConnectionsMap.has(ip)) {
    return { allowed: true };
  }

  const connectionCount = ipConnectionsMap.get(ip)!.connections.size;
  
  if (connectionCount >= maxConnectionsPerIP) {
    return {
      allowed: false,
      reason: `Maximum connections per IP (${maxConnectionsPerIP}) exceeded`
    };
  }

  return { allowed: true };
}

/**
 * Check if an IP exceeds connection rate limit
 *
 * @param ip - The IP address to check
 * @param ipConnectionsMap - Map of IPs to connection info
 * @param rateLimit - Maximum connections per minute
 * @returns Result with allowed status and reason if blocked
 */
export function checkConnectionRate(
  ip: string,
  ipConnectionsMap: Map<string, IIpConnectionInfo>,
  rateLimit: number
): IIpValidationResult {
  const now = Date.now();
  const minute = 60 * 1000;

  // Get or create connection info
  if (!ipConnectionsMap.has(ip)) {
    const info: IIpConnectionInfo = {
      connections: new Set(),
      timestamps: [now],
      ipVariants: normalizeIP(ip)
    };
    ipConnectionsMap.set(ip, info);
    return { allowed: true };
  }

  // Get timestamps and filter out entries older than 1 minute
  const info = ipConnectionsMap.get(ip)!;
  const timestamps = info.timestamps.filter(time => now - time < minute);
  timestamps.push(now);
  info.timestamps = timestamps;

  // Check if rate exceeds limit
  if (timestamps.length > rateLimit) {
    return {
      allowed: false,
      reason: `Connection rate limit (${rateLimit}/min) exceeded`
    };
  }

  return { allowed: true };
}

/**
 * Track a connection for an IP
 *
 * @param ip - The IP address
 * @param connectionId - The connection ID to track
 * @param ipConnectionsMap - Map of IPs to connection info
 */
export function trackConnection(
  ip: string,
  connectionId: string,
  ipConnectionsMap: Map<string, IIpConnectionInfo>
): void {
  if (!ipConnectionsMap.has(ip)) {
    ipConnectionsMap.set(ip, {
      connections: new Set([connectionId]),
      timestamps: [Date.now()],
      ipVariants: normalizeIP(ip)
    });
    return;
  }

  const info = ipConnectionsMap.get(ip)!;
  info.connections.add(connectionId);
}

/**
 * Remove connection tracking for an IP
 *
 * @param ip - The IP address
 * @param connectionId - The connection ID to remove
 * @param ipConnectionsMap - Map of IPs to connection info
 */
export function removeConnection(
  ip: string,
  connectionId: string,
  ipConnectionsMap: Map<string, IIpConnectionInfo>
): void {
  if (!ipConnectionsMap.has(ip)) return;

  const info = ipConnectionsMap.get(ip)!;
  info.connections.delete(connectionId);

  if (info.connections.size === 0) {
    ipConnectionsMap.delete(ip);
  }
}

/**
 * Clean up expired rate limits
 *
 * @param rateLimits - Map of rate limits to clean up
 * @param logger - Logger for debug messages
 */
export function cleanupExpiredRateLimits(
  rateLimits: Map<string, Map<string, IRateLimitInfo>>,
  logger?: ISecurityLogger
): void {
  const now = Date.now();
  let totalRemoved = 0;

  for (const [routeId, routeLimits] of rateLimits.entries()) {
    let removed = 0;
    for (const [key, limit] of routeLimits.entries()) {
      if (limit.expiry < now) {
        routeLimits.delete(key);
        removed++;
        totalRemoved++;
      }
    }
    
    if (removed > 0 && logger?.debug) {
      logger.debug(`Cleaned up ${removed} expired rate limits for route ${routeId}`);
    }
  }
  
  if (totalRemoved > 0 && logger?.info) {
    logger.info(`Cleaned up ${totalRemoved} expired rate limits total`);
  }
}

/**
 * Generate basic auth header value from username and password
 *
 * @param username - The username
 * @param password - The password
 * @returns Base64 encoded basic auth string
 */
export function generateBasicAuthHeader(username: string, password: string): string {
  return `Basic ${Buffer.from(`${username}:${password}`).toString('base64')}`;
}

/**
 * Parse basic auth header
 *
 * @param authHeader - The Authorization header value
 * @returns Username and password, or null if invalid
 */
export function parseBasicAuthHeader(
  authHeader: string
): { username: string; password: string } | null {
  if (!authHeader || !authHeader.startsWith('Basic ')) {
    return null;
  }

  try {
    const base64 = authHeader.slice(6); // Remove 'Basic '
    const decoded = Buffer.from(base64, 'base64').toString();
    const [username, password] = decoded.split(':');
    
    if (!username || !password) {
      return null;
    }

    return { username, password };
  } catch (err) {
    return null;
  }
}