#!/usr/bin/env node

import dotenv from 'dotenv';
dotenv.config({ path: '.env.local' });

import fs from 'fs';
import axios from 'axios';
import mysql from 'mysql2/promise';

const logFile = fs.createWriteStream('cve-sync.log', {
  flags: 'a',
  encoding: 'utf8',
});
const RESUME_FILE = '.enrichment_resume';
const DB = await mysql.createConnection({
  host: process.env.DB_HOST,
  port: process.env.DB_PORT || 3306,
  user: process.env.DB_USER,
  password: process.env.DB_PASSWORD,
  database: process.env.DB_NAME,
});
const BASE_URL = 'https://services.nvd.nist.gov/rest/json/cves/2.0';
const API_KEY = process.env.NVD_API_KEY;
const RESULTS_PER_PAGE = 2000;
let MAX_RANGE_DAYS = Number(process.env.NVD_MAX_RANGE_DAYS || 120);
log(`🧪 Getting CVEs from the last ${MAX_RANGE_DAYS} of days`);

if (MAX_RANGE_DAYS > 120) {
  log("⚠️ MAX_RANGE_DAYS exceeds NVD API limit. Defaulting to 120.");
  MAX_RANGE_DAYS = 120;
}
if (isNaN(MAX_RANGE_DAYS) || MAX_RANGE_DAYS < 1) {
  log("⚠️ Invalid MAX_RANGE_DAYS. Using default of 7.");
  MAX_RANGE_DAYS = 7;
}

function saveLastProcessedCVE(cveId) {
  fs.writeFileSync(RESUME_FILE, cveId, 'utf8');
}

function loadLastProcessedCVE() {
  if (!fs.existsSync(RESUME_FILE)) return null;
  return fs.readFileSync(RESUME_FILE, 'utf8');
}

function log(msg) {
  const now = new Date();

  // Generate the locale string with hour12 enabled (e.g. "14 Apr 2025, 08:14:42 AM")
  const raw = now.toLocaleString('en-AU', {
    day: '2-digit',
    month: 'short',
    year: 'numeric',
    hour: '2-digit',
    minute: '2-digit',
    second: '2-digit',
    hour12: true,
  });

  // Regex to convert only the AM/PM to lowercase
  const formatted = raw.replace(/\b(AM|PM)\b/, (match) => match.toLowerCase());

  const timestamp = `[${formatted}]`;
  const line = `${timestamp} ${msg}`;
  console.log(line);
  logFile.write(`${line}\n`);
}

function formatDate(isoString) {
  if (!isoString) return null;
  const date = new Date(isoString);
  return date.toISOString().slice(0, 19).replace('T', ' ');
}

function formatShortDate(isoString) {
  return new Date(isoString).toLocaleDateString('en-AU', {
    day: '2-digit',
    month: 'short',
    year: 'numeric',
  }).replace(/\b([A-Z])([a-z]+)\b/, (_, a, b) => a + b); // Capitalize only first letter of month
}

function extractCpeParts(cpe) {
  const parts = cpe.split(':');
  return {
    vendor: parts[3] || null,
    product: parts[4] || null,
    version: parts[5] || null
  };
}

function addDaysToISO(dateISO, days) {
  const date = new Date(dateISO);
  date.setDate(date.getDate() + days);
  return date.toISOString();
}

async function fetchCVEPage(startIndex, startDate, endDate, extraOptions = {}) {
  try {
    const res = await axios.get(BASE_URL, {
      params: {
        pubStartDate: startDate,
        pubEndDate: endDate,
        startIndex,
        resultsPerPage: RESULTS_PER_PAGE,
        ...extraOptions,
      },
      headers: API_KEY ? { apiKey: API_KEY } : {}
    });

    return res.data;
  } catch (err) {
    log(`❌ API error: ${err.response?.status} - ${err.response?.data?.message || err.message}`);
    throw err;
  }
}

async function processCVE(cveWrapper) {
  const cve = cveWrapper.cve;
  const cveId = cve.id;

  const title = cve.titles?.find(t => t.lang === 'en')?.title || '';
  const desc = cve.descriptions?.find(d => d.lang === 'en')?.value || '';
  const published = formatDate(cve.published);
  const modified = formatDate(cve.lastModified);

  // CVSSv2
  const metricV2 = cve.metrics?.cvssMetricV2?.[0];
  const severityV2 = metricV2?.cvssData?.baseSeverity || null;
  const scoreV2 = metricV2?.cvssData?.baseScore || null;
  const vectorV2 = metricV2?.cvssData?.vectorString || '';

  // CVSSv3
  const metricV3 = cve.metrics?.cvssMetricV31?.[0];
  const severityV3 = metricV3?.cvssData?.baseSeverity || null;
  const scoreV3 = metricV3?.cvssData?.baseScore || null;
  const vectorV3 = metricV3?.cvssData?.vectorString || '';

  // CVSSv4
  const metricV4 = cve.metrics?.cvssMetricV40?.[0] || cve.metrics?.cvssMetricV4?.[0];
  const severityV4 = metricV4?.cvssData?.baseSeverity || null;
  const scoreV4 = metricV4?.cvssData?.baseScore || null;
  const vectorV4 = metricV4?.cvssData?.vectorString || '';

  // CWE IDs
  const cweIds = (cve.weaknesses || [])
    .flatMap(w => w.description || [])
    .filter(desc => desc.lang === 'en')
    .map(desc => desc.value)
    .join(',');

  // References
  const references = (cve.references || [])
    .map(ref => ref.url)
    .join(',');

  // Tags
  const cveTags = cve.cveMetadata?.cveTags || [];
  const hasKev = cveTags.includes('Known_Exploited_Vulnerability');
  const hasCertNotes = cveTags.includes('CERT-VN');
  const hasCertAlerts = cveTags.includes('US-CERT-TA');

  try {
    await DB.execute(
      `INSERT INTO cves (
        id, title, description, published_date, last_modified_date,
        severity_v2, cvss_score_v2, cvss_vector_v2,
        severity_v3, cvss_score_v3, cvss_vector_v3,
        severity_v4, cvss_score_v4, cvss_vector_v4,
        cwe_ids, \`references\`, hasKev, hasCertNotes, hasCertAlerts, source
      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
      ON DUPLICATE KEY UPDATE
        last_modified_date = VALUES(last_modified_date),
        severity_v2 = IFNULL(severity_v2, VALUES(severity_v2)),
        cvss_score_v2 = IFNULL(cvss_score_v2, VALUES(cvss_score_v2)),
        cvss_vector_v2 = IFNULL(cvss_vector_v2, VALUES(cvss_vector_v2)),
        severity_v3 = IFNULL(severity_v3, VALUES(severity_v3)),
        cvss_score_v3 = IFNULL(cvss_score_v3, VALUES(cvss_score_v3)),
        cvss_vector_v3 = IFNULL(cvss_vector_v3, VALUES(cvss_vector_v3)),
        severity_v4 = IFNULL(severity_v4, VALUES(severity_v4)),
        cvss_score_v4 = IFNULL(cvss_score_v4, VALUES(cvss_score_v4)),
        cvss_vector_v4 = IFNULL(cvss_vector_v4, VALUES(cvss_vector_v4)),
        cwe_ids = IFNULL(cwe_ids, VALUES(cwe_ids)),
        \`references\` = IFNULL(\`references\`, VALUES(\`references\`)),
        hasKev = VALUES(hasKev),
        hasCertNotes = VALUES(hasCertNotes),
        hasCertAlerts = VALUES(hasCertAlerts),
        source = VALUES(source)
      `,
      [
        cveId, title, desc, published, modified,
        severityV2, scoreV2, vectorV2,
        severityV3, scoreV3, vectorV3,
        severityV4, scoreV4, vectorV4,
        cweIds, references, hasKev ? 1 : 0, hasCertNotes ? 1 : 0, hasCertAlerts ? 1 : 0, 'NVD'
      ]
    );
  } catch (err) {
    log(`❌ Error inserting CVE ${cveId}: ${err.message}`);
  }
}

async function getMostRecentModifiedDateFromDB() {
  const [rows] = await DB.query(`SELECT MAX(last_modified_date) AS lastMod FROM cves`);
  const lastMod = rows[0]?.lastMod;
  return lastMod ? new Date(lastMod).toISOString() : '2020-01-01T00:00:00.000Z';
}

async function importCVEFeed() {
  const now = new Date();
  const endDate = now.toISOString();

  const startDateObj = new Date(now);
  startDateObj.setDate(startDateObj.getDate() - MAX_RANGE_DAYS);
  const startDate = startDateObj.toISOString();

  log(`🚀 CVE sync started`);
  log(`🔄 Initializing script...`);
  log(`📍 Launching script`);
  log(`📅 Starting CVE sync from ${startDate} to ${endDate}`);

  const humanStart = formatShortDate(startDate);
  const humanEnd = formatShortDate(endDate);
  log(`📡 Fetching modified CVEs from ${humanStart} to ${humanEnd}...`);

  let startIndex = 0;
  let totalResults = Infinity;
  let pageCount = 0;

  do {
    const data = await fetchCVEPage(startIndex, startDate, endDate);
    const vulnerabilities = data.vulnerabilities || [];
    totalResults = data.totalResults ?? vulnerabilities.length;

    if (vulnerabilities.length === 0) {
      log(`⚠️ No CVEs returned at index ${startIndex}`);
      break;
    }

    log(`📄 Page ${++pageCount} — Processing ${vulnerabilities.length} CVEs from index ${startIndex} of ~${totalResults}`);

    for (const vuln of vulnerabilities) {
      await processCVE(vuln);
    }

    startIndex += RESULTS_PER_PAGE;
    await new Promise((r) => setTimeout(r, 6000));
  } while (startIndex < totalResults);

  log('✅ CVE import complete!');
  await DB.end();
  logFile.end();
}

async function importCVEFeedBackfill() {
  const now = new Date();
  const resumeFrom = loadLastSyncedDate();
  let startFrom = resumeFrom ? new Date(resumeFrom) : now;

  const MAX_RANGE_DAYS = 120;

  log(resumeFrom
    ? `🔁 Resuming CVE backfill from ${formatShortDate(startFrom.toISOString())}`
    : `⏮️ Starting CVE backfill from today (${formatShortDate(startFrom.toISOString())})`
  );

  while (true) {
    const end = new Date(startFrom);
    const start = new Date(startFrom);
    start.setDate(start.getDate() - MAX_RANGE_DAYS + 1); // 120-day window

    const startISO = start.toISOString();
    const endISO = end.toISOString();
    const humanRange = `${formatShortDate(startISO)} to ${formatShortDate(endISO)}`;

    log(`📡 Fetching published CVEs from ${humanRange}...`);

    let startIndex = 0;
    let totalResults = Infinity;
    let pageCount = 0;

    try {
      do {
        const data = await fetchCVEPage(startIndex, startISO, endISO);
        const vulnerabilities = data.vulnerabilities || [];
        totalResults = data.totalResults ?? vulnerabilities.length;

        if (vulnerabilities.length === 0) {
          log(`⚠️ No CVEs returned for ${humanRange} at index ${startIndex}`);
          break;
        }

        log(`📄 Page ${++pageCount} — ${vulnerabilities.length} CVEs from index ${startIndex}`);

        for (const vuln of vulnerabilities) {
          await processCVE(vuln);
        }

        startIndex += RESULTS_PER_PAGE;
        await new Promise((r) => setTimeout(r, 6000));
      } while (startIndex < totalResults);

      // Move the window backward
      saveLastSyncedDate(start.toISOString());
      startFrom = start;
    } catch (err) {
      log(`❌ Error during ${humanRange}: ${err.message}`);
      break;
    }

    if (start < new Date('2002-01-01')) {
      log(`🛑 Reached earliest supported CVE publication date — halting backfill.`);
      break;
    }
  }

  log('✅ CVE backfill complete!');
  await DB.end();
  logFile.end();
}

async function importCVEEnrichmentFromDB() {
  log(`🔍 Starting CVE enrichment from existing database records...`);

  const lastProcessed = loadLastProcessedCVE();
  if (lastProcessed) {
    log(`🔁 Resuming enrichment from after ${lastProcessed}`);
  }

  const [rows] = await DB.query(`
    SELECT id
    FROM cves
    WHERE severity_v3 IS NULL OR cvss_score_v3 IS NULL OR cvss_vector_v3 IS NULL
    ORDER BY id ASC
  `);

  if (rows.length === 0) {
    log(`✅ All CVEs seem enriched! Nothing to do.`);
    await DB.end();
    logFile.end();
    return;
  }

  let startIndex = 0;
  if (lastProcessed) {
    const resumeIndex = rows.findIndex(r => r.id === lastProcessed);
    if (resumeIndex !== -1 && resumeIndex + 1 < rows.length) {
      startIndex = resumeIndex + 1; // Start *after* last processed
    }
  }

  log(`📄 Found ${rows.length} CVEs needing enrichment, starting at index ${startIndex}.`);

  for (let i = startIndex; i < rows.length; i++) {
    const row = rows[i];
    const cveId = row.id;

    try {
      const res = await axios.get(BASE_URL, {
        params: { cveId },
        headers: API_KEY ? { apiKey: API_KEY } : {}
      });

      const vulnerabilities = res.data.vulnerabilities || [];
      if (vulnerabilities.length > 0) {
        await processCVE(vulnerabilities[0]);
        log(`✅ Enriched ${cveId}`);
      } else {
        log(`⚠️ No data found for ${cveId}`);
      }

      saveLastProcessedCVE(cveId); // ✨ Save after each successful CVE
      await new Promise(r => setTimeout(r, 6000)); // Sleep to respect rate limit
    } catch (err) {
      log(`❌ Failed to fetch/enrich ${cveId}: ${err.message}`);
    }
  }

  log('✅ CVE enrichment batch complete!');
  fs.unlinkSync(RESUME_FILE); // ✨ Cleanup resume file if done
  await DB.end();
  logFile.end();
}

async function importCVEEnrichmentFast() {
  log(`🚀 Starting FAST CVE enrichment by modified date...`);

  const now = new Date();
  const earliestDate = new Date('2002-01-01T00:00:00.000Z');
  let endDateObj = new Date(); // most recent
  const STEP_DAYS = 120;

  while (endDateObj > earliestDate) {
    const startDateObj = new Date(endDateObj);
    startDateObj.setDate(startDateObj.getDate() - STEP_DAYS);

    const startISO = startDateObj.toISOString();
    const endISO = endDateObj.toISOString();
    const humanRange = `${formatShortDate(startISO)} to ${formatShortDate(endISO)}`;

    log(`📅 Fetching modified CVEs from ${humanRange}`);

    let startIndex = 0;
    let totalResults = 0;
    let processedCount = 0;
    let pageCount = 0;

    try {
      const initial = await fetchCVEPage(0, startISO, endISO);
      totalResults = initial.totalResults || 0;

      if (totalResults === 0) {
        log(`⚠️ No CVEs to process in this window.`);
        endDateObj = startDateObj;
        continue;
      }

      log(`📦 Found ${totalResults} CVEs to enrich from ${humanRange}`);

      while (startIndex < totalResults) {
        const data = startIndex === 0 ? initial : await fetchCVEPage(startIndex, startISO, endISO);
        const vulnerabilities = data.vulnerabilities || [];

        log(`📄 Page ${++pageCount} — ${vulnerabilities.length} CVEs (Index ${startIndex})`);

        for (const vuln of vulnerabilities) {
          await processCVE(vuln);
          processedCount++;
          if (processedCount % 100 === 0 || processedCount === totalResults) {
            const pct = ((processedCount / totalResults) * 100).toFixed(1);
            log(`📊 Progress: ${processedCount}/${totalResults} CVEs (${pct}%)`);
          }
        }

        startIndex += RESULTS_PER_PAGE;
        await new Promise(r => setTimeout(r, 6000)); // API rate limit
      }

    } catch (err) {
      log(`❌ Error during enrichment for ${humanRange}: ${err.message}`);
    }

    endDateObj = startDateObj; // 🕒 step backward
  }

  log('✅ Full enrichment pass complete!');
  await DB.end();
  logFile.end();
}




importCVEEnrichmentFast().catch((err) => {
  log(`❌ Fatal error during enrichment: ${err.message}`);
  logFile.end();
});