Checkpoint/src/utils/bot-range-downloader.ts

import { TimedDownloadManager, type TimedDownloadSource } from './timed-downloads.js';
import { type DurationInput } from './time.js';
import { validateCIDR, isValidIP, ipToCIDR } from './ip-validation.js';
import * as logs from './logs.js';

// ==================== TYPE DEFINITIONS ====================

export interface BotSource {
  readonly name: string;
  readonly url: string;
  readonly updateInterval: DurationInput; // Uses time.ts format: "24h", "5m", etc.
  readonly dnsVerificationDomain?: string;
  readonly enabled: boolean;
}

export interface IPRange {
  readonly cidr: string;
  readonly ipv4?: boolean;
  readonly ipv6?: boolean;
}

export interface BotIPRanges {
  readonly botName: string;
  readonly ranges: readonly IPRange[];
  readonly lastUpdated: number;
  readonly source: string;
}

// ==================== UNIVERSAL PARSER ====================

/**
 * Universal parser that extracts IP ranges from any format and converts to CIDR list
 */
class UniversalRangeParser {
  static parse(data: string): readonly IPRange[] {
    const ranges: IPRange[] = [];
    const trimmed = data.trim();
    
    logs.plugin('bot-range-downloader', `Parsing ${trimmed.length} bytes of data`);
    
    // Try JSON parsing first
    let parsedFromJSON = false;
    try {
      const parsed = JSON.parse(trimmed);
      
      // Handle Google's JSON format: { "prefixes": [{"ipv4Prefix": "..."}, {"ipv6Prefix": "..."}] }
      if (parsed.prefixes && Array.isArray(parsed.prefixes)) {
        for (const prefix of parsed.prefixes) {
          if (prefix.ipv4Prefix) {
            const cidrResult = validateCIDR(prefix.ipv4Prefix);
            if (cidrResult.valid) {
              ranges.push({ 
                cidr: prefix.ipv4Prefix, 
                ipv4: cidrResult.type === 'ipv4', 
                ipv6: cidrResult.type !== 'ipv4' 
              });
            }
          }
          if (prefix.ipv6Prefix) {
            const cidrResult = validateCIDR(prefix.ipv6Prefix);
            if (cidrResult.valid) {
              ranges.push({ 
                cidr: prefix.ipv6Prefix, 
                ipv4: cidrResult.type === 'ipv4', 
                ipv6: cidrResult.type !== 'ipv4' 
              });
            }
          }
        }
        parsedFromJSON = true;
      }
      
      // Handle Microsoft/generic JSON format: { "ranges": ["...", "..."] }
      else if (parsed.ranges && Array.isArray(parsed.ranges)) {
        for (const range of parsed.ranges) {
          if (typeof range === 'string') {
            const cidrResult = validateCIDR(range);
            if (cidrResult.valid) {
              ranges.push({ 
                cidr: range, 
                ipv4: cidrResult.type === 'ipv4', 
                ipv6: cidrResult.type !== 'ipv4' 
              });
            }
          }
        }
        parsedFromJSON = true;
      }
      
      // Handle simple JSON array: ["...", "..."]
      else if (Array.isArray(parsed)) {
        for (const item of parsed) {
          if (typeof item === 'string') {
            // Check if it's already CIDR or needs conversion
            if (item.includes('/')) {
              const cidrResult = validateCIDR(item);
              if (cidrResult.valid) {
                ranges.push({ 
                  cidr: item, 
                  ipv4: cidrResult.type === 'ipv4', 
                  ipv6: cidrResult.type !== 'ipv4' 
                });
              }
            } else if (isValidIP(item)) {
              // Convert single IP to CIDR notation
              const cidr = ipToCIDR(item);
              if (cidr) {
                const cidrResult = validateCIDR(cidr);
                if (cidrResult.valid) {
                  ranges.push({ 
                    cidr, 
                    ipv4: cidrResult.type === 'ipv4', 
                    ipv6: cidrResult.type !== 'ipv4' 
                  });
                }
              }
            }
          }
        }
        parsedFromJSON = true;
      }
      
    } catch {
      // Not JSON, continue with text parsing
    }
    
    // If we successfully parsed JSON, return those results
    if (parsedFromJSON) {
      logs.plugin('bot-range-downloader', `Parsed ${ranges.length} ranges from JSON format`);
      return ranges.slice(0, 100000);
    }
    
    // Text-based parsing - handle both CIDR lists and IP lists
    const lines = trimmed.split('\n');
    
    for (const line of lines) {
      const cleaned = line.trim();
      
      // Skip empty lines and comments
      if (!cleaned || cleaned.startsWith('#') || cleaned.startsWith('//')) {
        continue;
      }
      
      // Check if line contains CIDR notation
      if (cleaned.includes('/')) {
        const cidrResult = validateCIDR(cleaned);
        if (cidrResult.valid) {
          ranges.push({ 
            cidr: cleaned, 
            ipv4: cidrResult.type === 'ipv4', 
            ipv6: cidrResult.type !== 'ipv4' 
          });
        }
      } 
      // Check if line is a single IP address
      else if (isValidIP(cleaned)) {
        // Convert single IP to CIDR notation
        const cidr = ipToCIDR(cleaned);
        if (cidr) {
          const cidrResult = validateCIDR(cidr);
          if (cidrResult.valid) {
            ranges.push({ 
              cidr, 
              ipv4: cidrResult.type === 'ipv4', 
              ipv6: cidrResult.type !== 'ipv4' 
            });
          }
        }
      }
    }
    
    logs.plugin('bot-range-downloader', `Parsed ${ranges.length} ranges from text format`);
    return ranges.slice(0, 100000);
  }
}

// ==================== BOT RANGE DOWNLOADER ====================

export class BotRangeDownloader {
  private readonly downloadManager: TimedDownloadManager;

  constructor() {
    this.downloadManager = new TimedDownloadManager('bot-ranges');
  }

  /**
   * Converts bot source to generic timed download source
   */
  private createTimedDownloadSource(botSource: BotSource): TimedDownloadSource {
    return {
      name: botSource.name,
      url: botSource.url,
      updateInterval: botSource.updateInterval,
      enabled: botSource.enabled,
      parser: {
        format: 'custom',
        parseFunction: (data: string) => {
          const ranges = UniversalRangeParser.parse(data);
          return {
            ranges,
            lastUpdated: Date.now(),
            source: botSource.url
          };
        },
      },
      validator: {
        maxSize: 50 * 1024 * 1024, // 50MB max
        maxEntries: 100000,
        validationFunction: (data: unknown): boolean => {
          return !!(data && typeof data === 'object' && 
                   'ranges' in data && Array.isArray((data as any).ranges) && 
                   (data as any).ranges.length > 0);
        },
      },
      headers: {
        'Accept': 'application/json, text/plain, */*',
        'User-Agent': 'Checkpoint-Security-Gateway/1.0 (Bot Range Downloader)',
      },
    };
  }

  /**
   * Downloads bot ranges using the universal parser
   */
  async downloadBotRanges(botSource: BotSource): Promise<{ success: boolean; ranges?: readonly IPRange[]; error?: string }> {
    const timedSource = this.createTimedDownloadSource(botSource);
    const result = await this.downloadManager.downloadFromSource(timedSource);
    
    if (result.success && result.data) {
      const parsedData = result.data as { ranges: readonly IPRange[] };
      return {
        success: true,
        ranges: parsedData.ranges,
      };
    } else {
      return {
        success: false,
        error: result.error,
      };
    }
  }

  /**
   * Loads bot ranges from disk
   */
  async loadBotRanges(botName: string): Promise<BotIPRanges | null> {
    const downloadedData = await this.downloadManager.loadDownloadedData(botName);
    if (!downloadedData) {
      return null;
    }

    const data = downloadedData.data as { ranges: readonly IPRange[] };
    
    return {
      botName,
      ranges: data.ranges,
      lastUpdated: downloadedData.lastUpdated,
      source: downloadedData.source,
    };
  }

  /**
   * Checks if bot ranges need updating
   */
  async needsUpdate(botSource: BotSource): Promise<boolean> {
    const timedSource = this.createTimedDownloadSource(botSource);
    return await this.downloadManager.needsUpdate(timedSource);
  }

  /**
   * Starts periodic updates for bot sources
   */
  startPeriodicUpdates(botSources: readonly BotSource[]): void {
    const timedSources = botSources.map(source => this.createTimedDownloadSource(source));
    this.downloadManager.startPeriodicUpdates(timedSources);
  }

  /**
   * Updates all bot sources that need updating
   */
  async updateAllSources(botSources: readonly BotSource[]): Promise<void> {
    const timedSources = botSources.map(source => this.createTimedDownloadSource(source));
    await this.downloadManager.updateAllSources(timedSources);
  }
}