smartmta/ts/security/classes.ipreputationchecker.ts

import * as plugins from '../plugins.js';
import * as paths from '../paths.js';
import { logger } from '../logger.js';
import { SecurityLogger, SecurityLogLevel, SecurityEventType } from './classes.securitylogger.js';
import { RustSecurityBridge } from './classes.rustsecuritybridge.js';
import { LRUCache } from 'lru-cache';

/**
 * Reputation check result information
 */
export interface IReputationResult {
  score: number;          // 0 (worst) to 100 (best)
  isSpam: boolean;        // true if the IP is known for spam
  isProxy: boolean;       // true if the IP is a known proxy
  isTor: boolean;         // true if the IP is a known Tor exit node
  isVPN: boolean;         // true if the IP is a known VPN
  country?: string;       // Country code (if available)
  asn?: string;           // Autonomous System Number (if available)
  org?: string;           // Organization name (if available)
  blacklists?: string[];  // Names of blacklists that include this IP
  timestamp: number;      // When this result was created/retrieved
  error?: string;         // Error message if check failed
}

/**
 * Reputation threshold scores
 */
export enum ReputationThreshold {
  HIGH_RISK = 20,      // Score below this is considered high risk
  MEDIUM_RISK = 50,    // Score below this is considered medium risk
  LOW_RISK = 80        // Score below this is considered low risk (but not trusted)
}

/**
 * IP type classifications
 */
export enum IPType {
  RESIDENTIAL = 'residential',
  DATACENTER = 'datacenter',
  PROXY = 'proxy',
  TOR = 'tor',
  VPN = 'vpn',
  UNKNOWN = 'unknown'
}

/**
 * Options for the IP Reputation Checker
 */
export interface IIPReputationOptions {
  maxCacheSize?: number;  // Maximum number of IPs to cache (default: 10000)
  cacheTTL?: number;      // TTL for cache entries in ms (default: 24 hours)
  dnsblServers?: string[]; // List of DNSBL servers to check
  highRiskThreshold?: number; // Score below this is high risk
  mediumRiskThreshold?: number; // Score below this is medium risk
  lowRiskThreshold?: number; // Score below this is low risk
  enableLocalCache?: boolean; // Whether to persist cache to disk (default: true)
  enableDNSBL?: boolean; // Whether to use DNSBL checks (default: true)
  enableIPInfo?: boolean; // Whether to use IP info service (default: true)
}

/**
 * IP reputation checker — delegates DNSBL lookups to the Rust security bridge.
 * Retains LRU caching and disk persistence in TypeScript.
 */
export class IPReputationChecker {
  private static instance: IPReputationChecker;
  private reputationCache: LRUCache<string, IReputationResult>;
  private options: Required<IIPReputationOptions>;
  private storageManager?: any;

  private static readonly DEFAULT_OPTIONS: Required<IIPReputationOptions> = {
    maxCacheSize: 10000,
    cacheTTL: 24 * 60 * 60 * 1000,
    dnsblServers: [],
    highRiskThreshold: ReputationThreshold.HIGH_RISK,
    mediumRiskThreshold: ReputationThreshold.MEDIUM_RISK,
    lowRiskThreshold: ReputationThreshold.LOW_RISK,
    enableLocalCache: true,
    enableDNSBL: true,
    enableIPInfo: true
  };

  constructor(options: IIPReputationOptions = {}, storageManager?: any) {
    this.options = {
      ...IPReputationChecker.DEFAULT_OPTIONS,
      ...options
    };

    this.storageManager = storageManager;

    this.reputationCache = new LRUCache<string, IReputationResult>({
      max: this.options.maxCacheSize,
      ttl: this.options.cacheTTL,
    });

    if (this.options.enableLocalCache) {
      this.loadCache().catch(error => {
        logger.log('error', `Failed to load IP reputation cache during initialization: ${error.message}`);
      });
    }
  }

  public static getInstance(options: IIPReputationOptions = {}, storageManager?: any): IPReputationChecker {
    if (!IPReputationChecker.instance) {
      IPReputationChecker.instance = new IPReputationChecker(options, storageManager);
    }
    return IPReputationChecker.instance;
  }

  /**
   * Check an IP address's reputation via the Rust bridge
   */
  public async checkReputation(ip: string): Promise<IReputationResult> {
    try {
      if (!this.isValidIPAddress(ip)) {
        logger.log('warn', `Invalid IP address format: ${ip}`);
        return this.createErrorResult(ip, 'Invalid IP address format');
      }

      // Check cache first
      const cachedResult = this.reputationCache.get(ip);
      if (cachedResult) {
        logger.log('info', `Using cached reputation data for IP ${ip}`, {
          score: cachedResult.score,
          isSpam: cachedResult.isSpam
        });
        return cachedResult;
      }

      // Delegate to Rust bridge
      const bridge = RustSecurityBridge.getInstance();
      const rustResult = await bridge.checkIpReputation(ip);

      const result: IReputationResult = {
        score: rustResult.score,
        isSpam: rustResult.listed_count > 0,
        isProxy: rustResult.ip_type === 'proxy',
        isTor: rustResult.ip_type === 'tor',
        isVPN: rustResult.ip_type === 'vpn',
        blacklists: rustResult.dnsbl_results
          .filter(d => d.listed)
          .map(d => d.server),
        timestamp: Date.now(),
      };

      this.reputationCache.set(ip, result);

      if (this.options.enableLocalCache) {
        this.saveCache().catch(error => {
          logger.log('error', `Failed to save IP reputation cache: ${error.message}`);
        });
      }

      this.logReputationCheck(ip, result);
      return result;
    } catch (error) {
      logger.log('error', `Error checking IP reputation for ${ip}: ${error.message}`, {
        ip,
        stack: error.stack
      });
      const errorResult = this.createErrorResult(ip, error.message);
      // Cache error results to avoid repeated failing lookups
      this.reputationCache.set(ip, errorResult);
      return errorResult;
    }
  }

  private createErrorResult(ip: string, errorMessage: string): IReputationResult {
    return {
      score: 50,
      isSpam: false,
      isProxy: false,
      isTor: false,
      isVPN: false,
      timestamp: Date.now(),
      error: errorMessage
    };
  }

  private isValidIPAddress(ip: string): boolean {
    const ipv4Pattern = /^((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$/;
    return ipv4Pattern.test(ip);
  }

  private logReputationCheck(ip: string, result: IReputationResult): void {
    let logLevel = SecurityLogLevel.INFO;
    if (result.score < this.options.highRiskThreshold) {
      logLevel = SecurityLogLevel.WARN;
    }

    SecurityLogger.getInstance().logEvent({
      level: logLevel,
      type: SecurityEventType.IP_REPUTATION,
      message: `IP reputation check ${result.isSpam ? 'flagged spam' : 'completed'} for ${ip}`,
      ipAddress: ip,
      details: {
        score: result.score,
        isSpam: result.isSpam,
        isProxy: result.isProxy,
        isTor: result.isTor,
        isVPN: result.isVPN,
        country: result.country,
        blacklists: result.blacklists
      },
      success: !result.isSpam
    });
  }

  private async saveCache(): Promise<void> {
    try {
      const entries = Array.from(this.reputationCache.entries()).map(([ip, data]) => ({
        ip,
        data
      }));

      if (entries.length === 0) {
        return;
      }

      const cacheData = JSON.stringify(entries);

      if (this.storageManager) {
        await this.storageManager.set('/security/ip-reputation-cache.json', cacheData);
        logger.log('info', `Saved ${entries.length} IP reputation cache entries to StorageManager`);
      } else {
        const cacheDir = plugins.path.join(paths.dataDir, 'security');
        await plugins.smartfs.directory(cacheDir).recursive().create();
        const cacheFile = plugins.path.join(cacheDir, 'ip_reputation_cache.json');
        await plugins.smartfs.file(cacheFile).write(cacheData);
        logger.log('info', `Saved ${entries.length} IP reputation cache entries to disk`);
      }
    } catch (error) {
      logger.log('error', `Failed to save IP reputation cache: ${error.message}`);
    }
  }

  private async loadCache(): Promise<void> {
    try {
      let cacheData: string | null = null;
      let fromFilesystem = false;

      if (this.storageManager) {
        try {
          cacheData = await this.storageManager.get('/security/ip-reputation-cache.json');

          if (!cacheData) {
            const cacheFile = plugins.path.join(paths.dataDir, 'security', 'ip_reputation_cache.json');
            if (plugins.fs.existsSync(cacheFile)) {
              logger.log('info', 'Migrating IP reputation cache from filesystem to StorageManager');
              cacheData = plugins.fs.readFileSync(cacheFile, 'utf8');
              fromFilesystem = true;
              await this.storageManager.set('/security/ip-reputation-cache.json', cacheData);
              logger.log('info', 'IP reputation cache migrated to StorageManager successfully');
              try {
                plugins.fs.unlinkSync(cacheFile);
                logger.log('info', 'Old cache file removed after migration');
              } catch (deleteError) {
                logger.log('warn', `Could not delete old cache file: ${deleteError.message}`);
              }
            }
          }
        } catch (error) {
          logger.log('error', `Error loading from StorageManager: ${error.message}`);
        }
      } else {
        const cacheFile = plugins.path.join(paths.dataDir, 'security', 'ip_reputation_cache.json');
        if (plugins.fs.existsSync(cacheFile)) {
          cacheData = plugins.fs.readFileSync(cacheFile, 'utf8');
          fromFilesystem = true;
        }
      }

      if (cacheData) {
        const entries = JSON.parse(cacheData);
        const now = Date.now();
        const validEntries = entries.filter(entry => {
          const age = now - entry.data.timestamp;
          return age < this.options.cacheTTL;
        });

        for (const entry of validEntries) {
          this.reputationCache.set(entry.ip, entry.data);
        }

        const source = fromFilesystem ? 'disk' : 'StorageManager';
        logger.log('info', `Loaded ${validEntries.length} IP reputation cache entries from ${source}`);
      }
    } catch (error) {
      logger.log('error', `Failed to load IP reputation cache: ${error.message}`);
    }
  }

  public static getRiskLevel(score: number): 'high' | 'medium' | 'low' | 'trusted' {
    if (score < ReputationThreshold.HIGH_RISK) {
      return 'high';
    } else if (score < ReputationThreshold.MEDIUM_RISK) {
      return 'medium';
    } else if (score < ReputationThreshold.LOW_RISK) {
      return 'low';
    } else {
      return 'trusted';
    }
  }

  public updateStorageManager(storageManager: any): void {
    this.storageManager = storageManager;
    logger.log('info', 'IPReputationChecker storage manager updated');

    if (this.options.enableLocalCache && this.reputationCache.size > 0) {
      this.saveCache().catch(error => {
        logger.log('error', `Failed to save cache to new storage manager: ${error.message}`);
      });
    }
  }
}