Checkpoint/checkpoint.js

import { registerPlugin, loadConfig, rootDir } from './index.js';
import crypto from 'crypto';
import path from 'path';
import fs from 'fs';
import { promises as fsPromises } from 'fs';
import { dirname, join } from 'path';
import { fileURLToPath } from 'url';
import { Level } from 'level';
import cookie from 'cookie';
import { parseDuration } from './utils/time.js';
import { getRealIP } from './utils/network.js';
import ttl from 'level-ttl';
import { Readable } from 'stream';
import {
  challengeStore,
  generateRequestID as proofGenerateRequestID,
  getChallengeParams,
  deleteChallenge,
  verifyPoW,
  verifyPoS,
} from './utils/proof.js';
import express from 'express';
// Import recordEvent dynamically to avoid circular dependency issues
let recordEvent;
let statsLoadPromise = import('./plugins/stats.js')
  .then((stats) => {
    recordEvent = stats.recordEvent;
  })
  .catch((err) => {
    console.error('Failed to import stats module:', err);
    recordEvent = null;
  });

function sanitizePath(inputPath) {
  let pathOnly = inputPath.replace(/[\x00-\x1F\x7F]/g, '');

  pathOnly = pathOnly.replace(/[<>;"'`|]/g, '');

  const parts = pathOnly.split('/').filter((seg) => seg && seg !== '.' && seg !== '..');

  return '/' + parts.map((seg) => encodeURIComponent(seg)).join('/');
}

const checkpointConfig = {};
let hmacSecret = null;
const usedNonces = new Map();
const ipRateLimit = new Map();

const tokenCache = new Map();

let db;

const tokenExpirations = new Map();

let interstitialTemplate = null;

const __dirname = dirname(fileURLToPath(import.meta.url));

function simpleTemplate(str) {
  return function (data) {
    return str.replace(/\{\{\s*([^{}]+?)\s*\}\}/g, (_, key) => {
      let value = data;

      for (const part of key.trim().split('.')) {
        value = value?.[part];
        if (value == null) break;
      }
      return value != null ? String(value) : '';
    });
  };
}

async function initConfig() {
  await loadConfig('checkpoint', checkpointConfig);

  // Handle new nested configuration structure
  // Map nested structure to flat structure for internal use
  checkpointConfig.Enabled = checkpointConfig.Core.Enabled;
  checkpointConfig.CookieName = checkpointConfig.Core.CookieName;
  checkpointConfig.CookieDomain = checkpointConfig.Core.CookieDomain;
  checkpointConfig.SanitizeURLs = checkpointConfig.Core.SanitizeURLs;

  // Proof of Work settings
  checkpointConfig.Difficulty = checkpointConfig.ProofOfWork.Difficulty;
  checkpointConfig.SaltLength = checkpointConfig.ProofOfWork.SaltLength;
  checkpointConfig.ChallengeExpiration = parseDuration(
    checkpointConfig.ProofOfWork.ChallengeExpiration,
  );
  checkpointConfig.MaxAttemptsPerHour = checkpointConfig.ProofOfWork.MaxAttemptsPerHour;

  // Proof of Space-Time settings
  checkpointConfig.CheckPoSTimes = checkpointConfig.ProofOfSpaceTime.Enabled;
  checkpointConfig.PoSTimeConsistencyRatio = checkpointConfig.ProofOfSpaceTime.ConsistencyRatio;

  // Token settings
  checkpointConfig.TokenExpiration = parseDuration(checkpointConfig.Token.Expiration);
  checkpointConfig.MaxNonceAge = parseDuration(checkpointConfig.Token.MaxNonceAge);

  // Storage settings
  checkpointConfig.SecretConfigPath = checkpointConfig.Storage.SecretPath;
  checkpointConfig.TokenStoreDBPath = checkpointConfig.Storage.TokenDBPath;
  checkpointConfig.InterstitialPaths = checkpointConfig.Storage.InterstitialTemplates;

  // Process exclusions
  checkpointConfig.ExclusionRules = checkpointConfig.Exclusion || [];

  // Process bypass keys
  checkpointConfig.BypassQueryKeys = [];
  checkpointConfig.BypassHeaderKeys = [];
  checkpointConfig.BypassKeys.forEach((key) => {
    if (key.Type === 'query') {
      checkpointConfig.BypassQueryKeys.push({
        Key: key.Key,
        Value: key.Value,
        Domains: key.Hosts || [],
      });
    } else if (key.Type === 'header') {
      checkpointConfig.BypassHeaderKeys.push({
        Name: key.Key,
        Value: key.Value,
        Domains: key.Hosts || [],
      });
    }
  });

  // Extension handling
  checkpointConfig.HTMLCheckpointIncludedExtensions = checkpointConfig.Extensions?.IncludeOnly || [];
  checkpointConfig.HTMLCheckpointExcludedExtensions = checkpointConfig.Extensions?.Exclude || [];

  // Remove legacy arrays
  checkpointConfig.HTMLCheckpointExclusions = [];
  checkpointConfig.UserAgentValidationExclusions = [];
  checkpointConfig.UserAgentRequiredPrefixes = {};
  checkpointConfig.ReverseProxyMappings = {};
}

function addReadStreamSupport(dbInstance) {
  if (!dbInstance.createReadStream) {
    dbInstance.createReadStream = (opts) =>
      Readable.from(
        (async function* () {
          for await (const [key, value] of dbInstance.iterator(opts)) {
            yield { key, value };
          }
        })(),
      );
  }
  return dbInstance;
}

function initTokenStore() {
  try {
    const storePath = join(rootDir, checkpointConfig.TokenStoreDBPath || 'db/tokenstore');
    fs.mkdirSync(storePath, { recursive: true });

    let rawDB = new Level(storePath, { valueEncoding: 'json' });

    addReadStreamSupport(rawDB);
    db = ttl(rawDB, { defaultTTL: checkpointConfig.TokenExpiration });

    addReadStreamSupport(db);
    console.log('Token store initialized with TTL');
  } catch (err) {
    console.error('Failed to initialize token store:', err);
  }
}

function getFullClientIP(request) {
  const ip = getRealIP(request) || '';
  const h = crypto.createHash('sha256').update(ip).digest();
  return h.slice(0, 8).toString('hex');
}

function hashUserAgent(ua) {
  if (!ua) return '';
  const h = crypto.createHash('sha256').update(ua).digest();
  return h.slice(0, 8).toString('hex');
}

function extractBrowserFingerprint(request) {
  const headers = [
    'sec-ch-ua',
    'sec-ch-ua-platform',
    'sec-ch-ua-mobile',
    'sec-ch-ua-platform-version',
    'sec-ch-ua-arch',
    'sec-ch-ua-model',
  ];
  const parts = headers.map((h) => request.headers.get(h)).filter(Boolean);
  if (!parts.length) return '';
  const buf = Buffer.from(parts.join('|'));
  const h = crypto.createHash('sha256').update(buf).digest();
  return h.slice(0, 12).toString('hex');
}

async function getInterstitialTemplate() {
  if (!interstitialTemplate) {
    for (const p of checkpointConfig.InterstitialPaths) {
      try {
        let templatePath = join(__dirname, p);
        if (fs.existsSync(templatePath)) {
          const raw = await fsPromises.readFile(templatePath, 'utf8');
          interstitialTemplate = simpleTemplate(raw);
          break;
        }

        templatePath = join(rootDir, p);
        if (fs.existsSync(templatePath)) {
          const raw = await fsPromises.readFile(templatePath, 'utf8');
          interstitialTemplate = simpleTemplate(raw);
          break;
        }
      } catch (e) {
        console.warn(`Failed to load interstitial template from path ${p}:`, e);
      }
    }

    if (!interstitialTemplate) {
      // Create a minimal fallback template
      console.warn('Could not find interstitial HTML template, using minimal fallback');
      interstitialTemplate = simpleTemplate(`
<!DOCTYPE html>
<html>
<head>
  <title>Security Verification</title>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
</head>
<body>
  <h1>Security Verification Required</h1>
  <p>Please wait while we verify your request...</p>
  <div id="verification-data" 
       data-target="{{TargetPath}}" 
       data-request-id="{{RequestID}}">
  </div>
  <script src="/js/c.js"></script>
</body>
</html>
      `);
    }
  }
  return interstitialTemplate;
}

// Helper function for safe stats recording
function safeRecordEvent(metric, data) {
  // If recordEvent is not yet loaded, try to wait for it
  if (!recordEvent && statsLoadPromise) {
    statsLoadPromise.then(() => {
      if (recordEvent) {
        try {
          recordEvent(metric, data);
        } catch (err) {
          console.error(`Failed to record ${metric} event:`, err);
        }
      }
    });
    return;
  }

  if (typeof recordEvent === 'function') {
    try {
      recordEvent(metric, data);
    } catch (err) {
      console.error(`Failed to record ${metric} event:`, err);
    }
  }
}

async function serveInterstitial(request) {
  const ip = getRealIP(request);
  const requestPath = new URL(request.url).pathname;
  safeRecordEvent('checkpoint.sent', { ip, path: requestPath });
  let tpl;
  try {
    tpl = await getInterstitialTemplate();
  } catch (err) {
    console.error('Interstitial template error:', err);
    return new Response('Security verification required.', {
      status: 200,
      headers: { 'Content-Type': 'text/plain' },
    });
  }

  const requestID = proofGenerateRequestID(request, checkpointConfig);
  const url = new URL(request.url);
  const host = request.headers.get('host') || url.hostname;
  const targetPath = url.pathname;
  const fullURL = request.url;

  const html = tpl({
    TargetPath: targetPath,
    RequestID: requestID,
    Host: host,
    FullURL: fullURL,
  });

  return new Response(html, {
    status: 200,
    headers: { 'Content-Type': 'text/html; charset=utf-8' },
  });
}

async function handleGetCheckpointChallenge(request) {
  const url = new URL(request.url);
  const requestID = url.searchParams.get('id');
  if (!requestID) {
    return new Response(JSON.stringify({ error: 'Missing request ID' }), {
      status: 400,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  const ip = getRealIP(request);
  const attempts = (ipRateLimit.get(ip) || 0) + 1;
  ipRateLimit.set(ip, attempts);

  if (attempts > checkpointConfig.MaxAttemptsPerHour) {
    return new Response(
      JSON.stringify({ error: 'Too many challenge requests. Try again later.' }),
      {
        status: 429,
        headers: { 'Content-Type': 'application/json' },
      },
    );
  }

  const params = getChallengeParams(requestID);
  if (!params) {
    return new Response(JSON.stringify({ error: 'Challenge not found or expired' }), {
      status: 404,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  if (ip !== params.ClientIP) {
    return new Response(JSON.stringify({ error: 'IP address mismatch for challenge' }), {
      status: 403,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  const payload = {
    a: params.Challenge,
    b: params.Salt,
    c: params.Difficulty,
    d: params.PoSSeed,
  };
  return new Response(JSON.stringify(payload), {
    status: 200,
    headers: { 'Content-Type': 'application/json' },
  });
}

function calculateTokenHash(token) {
  const data = `${token.Nonce}:${token.Entropy}:${token.Created.getTime()}`;
  return crypto.createHash('sha256').update(data).digest('hex');
}

function computeTokenSignature(token) {
  const copy = { ...token, Signature: '' };
  const serialized = JSON.stringify(copy);
  return crypto.createHmac('sha256', hmacSecret).update(serialized).digest('hex');
}

function verifyTokenSignature(token) {
  if (!token.Signature) return false;
  const expected = computeTokenSignature(token);
  try {
    return crypto.timingSafeEqual(
      Buffer.from(token.Signature, 'hex'),
      Buffer.from(expected, 'hex'),
    );
  } catch (e) {
    return false;
  }
}

async function issueToken(request, token) {
  const tokenHash = calculateTokenHash(token);
  const storedData = {
    ClientIPHash: token.ClientIP,
    UserAgentHash: token.UserAgent,
    BrowserHint: token.BrowserHint,
    LastVerified: new Date(token.LastVerified).toISOString(),
    ExpiresAt: new Date(token.ExpiresAt).toISOString(),
  };

  try {
    await addToken(tokenHash, storedData);
  } catch (err) {
    console.error('Failed to store token:', err);
  }

  token.Signature = computeTokenSignature(token);

  const tokenStr = Buffer.from(JSON.stringify(token)).toString('base64');

  const url = new URL(request.url);
  const cookieDomain = checkpointConfig.CookieDomain || '';
  const sameSite = cookieDomain ? 'Lax' : 'Strict';
  const secure = url.protocol === 'https:';
  const expires = new Date(token.ExpiresAt).toUTCString();

  const domainPart = cookieDomain ? `; Domain=${cookieDomain}` : '';
  const securePart = secure ? '; Secure' : '';
  const cookieStr =
    `${checkpointConfig.CookieName}=${tokenStr}; Path=/` +
    `${domainPart}; Expires=${expires}; HttpOnly; SameSite=${sameSite}${securePart}`;
  return new Response(JSON.stringify({ token: tokenStr, expires_at: token.ExpiresAt }), {
    status: 200,
    headers: {
      'Content-Type': 'application/json',
      'Set-Cookie': cookieStr,
    },
  });
}

async function handleVerifyCheckpoint(request) {
  let body;
  try {
    body = await request.json();
  } catch (e) {
    safeRecordEvent('checkpoint.failure', { reason: 'invalid_json', ip: getRealIP(request) });
    return new Response(JSON.stringify({ error: 'Invalid JSON' }), {
      status: 400,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  const ip = getRealIP(request);
  const params = getChallengeParams(body.request_id);

  if (!params) {
    safeRecordEvent('checkpoint.failure', { reason: 'invalid_or_expired_request', ip });
    return new Response(JSON.stringify({ error: 'Invalid or expired request ID' }), {
      status: 400,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  if (ip !== params.ClientIP) {
    safeRecordEvent('checkpoint.failure', { reason: 'ip_mismatch', ip });
    return new Response(JSON.stringify({ error: 'IP address mismatch' }), {
      status: 403,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  const challenge = params.Challenge;
  const salt = params.Salt;

  if (!body.g || !verifyPoW(challenge, salt, body.g, params.Difficulty)) {
    safeRecordEvent('checkpoint.failure', { reason: 'invalid_pow', ip });
    return new Response(JSON.stringify({ error: 'Invalid proof-of-work solution' }), {
      status: 400,
      headers: { 'Content-Type': 'application/json' },
    });
  }

  const nonceKey = body.g + challenge;
  usedNonces.set(nonceKey, Date.now());

  if (body.h?.length === 3 && body.i?.length === 3) {
    try {
      verifyPoS(body.h, body.i, checkpointConfig);
    } catch (e) {
      safeRecordEvent('checkpoint.failure', { reason: 'invalid_pos', ip });
      return new Response(JSON.stringify({ error: e.message }), {
        status: 400,
        headers: { 'Content-Type': 'application/json' },
      });
    }
  }

  deleteChallenge(body.request_id);
  safeRecordEvent('checkpoint.success', { ip });
  const now = new Date();
  const expiresAt = new Date(now.getTime() + checkpointConfig.TokenExpiration);

  const token = {
    Nonce: body.g,
    ExpiresAt: expiresAt,
    ClientIP: getFullClientIP(request),
    UserAgent: hashUserAgent(request.headers.get('user-agent')),
    BrowserHint: extractBrowserFingerprint(request),
    Entropy: crypto.randomBytes(8).toString('hex'),
    Created: now,
    LastVerified: now,
    TokenFormat: 2,
  };

  token.Signature = computeTokenSignature(token);
  const tokenStr = Buffer.from(JSON.stringify(token)).toString('base64');

  const tokenKey = crypto.createHash('sha256').update(tokenStr).digest('hex');
  try {
    await db.put(tokenKey, true);
    tokenCache.set(tokenKey, true);

    tokenExpirations.set(tokenKey, new Date(token.ExpiresAt).getTime());
    console.log(`checkpoint: token stored in DB and cache key=${tokenKey}`);
  } catch (e) {
    console.error('checkpoint: failed to store token in DB:', e);
  }

  return new Response(JSON.stringify({ token: tokenStr, expires_at: token.ExpiresAt }), {
    status: 200,
    headers: { 'Content-Type': 'application/json' },
  });
}

function generateUpdatedCookie(token, secure) {
  token.Signature = computeTokenSignature(token);
  const tokenStr = Buffer.from(JSON.stringify(token)).toString('base64');
  const cookieDomain = checkpointConfig.CookieDomain || '';
  const sameSite = cookieDomain ? 'Lax' : 'Strict';
  const expires = new Date(token.ExpiresAt).toUTCString();

  const domainPart = cookieDomain ? `; Domain=${cookieDomain}` : '';
  const securePart = secure ? '; Secure' : '';
  const cookieStr =
    `${checkpointConfig.CookieName}=${tokenStr}; Path=/` +
    `${domainPart}; Expires=${expires}; HttpOnly; SameSite=${sameSite}${securePart}`;
  return cookieStr;
}

async function validateToken(tokenStr, request) {
  if (!tokenStr) return false;
  let token;
  try {
    token = JSON.parse(Buffer.from(tokenStr, 'base64').toString());
  } catch {
    console.log('checkpoint: invalid token format');
    return false;
  }

  if (Date.now() > new Date(token.ExpiresAt).getTime()) {
    console.log('checkpoint: token expired');
    return false;
  }

  if (!verifyTokenSignature(token)) {
    console.log('checkpoint: invalid token signature');
    return false;
  }

  const tokenKey = crypto.createHash('sha256').update(tokenStr).digest('hex');

  if (tokenCache.has(tokenKey)) return true;

  try {
    await db.get(tokenKey);
    tokenCache.set(tokenKey, true);

    tokenExpirations.set(tokenKey, new Date(token.ExpiresAt).getTime());
    return true;
  } catch {
    console.log('checkpoint: token not found in DB');
    return false;
  }
}

async function handleTokenRedirect(request) {
  const url = new URL(request.url);
  const tokenStr = url.searchParams.get('token');
  if (!tokenStr) return undefined;

  let token;
  try {
    token = JSON.parse(Buffer.from(tokenStr, 'base64').toString());

    if (Date.now() > new Date(token.ExpiresAt).getTime()) {
      console.log('checkpoint: token in URL parameter expired');
      return undefined;
    }

    if (!verifyTokenSignature(token)) {
      console.log('checkpoint: invalid token signature in URL parameter');
      return undefined;
    }

    const tokenKey = crypto.createHash('sha256').update(tokenStr).digest('hex');
    try {
      await db.get(tokenKey);
    } catch {
      console.log('checkpoint: token in URL parameter not found in DB');
      return undefined;
    }
  } catch (e) {
    console.log('checkpoint: invalid token format in URL parameter', e);
    return undefined;
  }

  const expires = new Date(token.ExpiresAt).toUTCString();
  const cookieDomain = checkpointConfig.CookieDomain || '';
  const sameSite = cookieDomain ? 'Lax' : 'Strict';
  const securePart = url.protocol === 'https:' ? '; Secure' : '';
  const domainPart = cookieDomain ? `; Domain=${cookieDomain}` : '';
  const cookieStr =
    `${checkpointConfig.CookieName}=${tokenStr}; Path=/` +
    `${domainPart}; Expires=${expires}; HttpOnly; SameSite=${sameSite}${securePart}`;

  url.searchParams.delete('token');
  const cleanUrl = url.pathname + (url.search || '');
  return new Response(null, {
    status: 302,
    headers: {
      'Set-Cookie': cookieStr,
      Location: cleanUrl,
    },
  });
}

function CheckpointMiddleware() {
  // Return Express-compatible middleware
  return {
    middleware: [
      // Add body parser middleware for JSON
      express.json({ limit: '10mb' }),
      // Main checkpoint middleware
      async (req, res, next) => {
        // Check if checkpoint is enabled
        if (checkpointConfig.Enabled === false) {
          return next();
        }

        // Convert Express request to the format expected by checkpoint logic
        const request = {
          url: `${req.protocol}://${req.get('host')}${req.originalUrl}`,
          method: req.method,
          headers: {
            get: (name) => req.get(name),
            entries: () => Object.entries(req.headers).map(([k, v]) => [k, Array.isArray(v) ? v.join(', ') : v])
          },
          json: () => Promise.resolve(req.body)
        };

        const urlObj = new URL(request.url);
        const host = request.headers.get('host')?.split(':')[0];
        const userAgent = request.headers.get('user-agent') || '';

        // 1) Bypass via query keys
        for (const { Key, Value, Domains } of checkpointConfig.BypassQueryKeys) {
          if (urlObj.searchParams.get(Key) === Value) {
            if (!Array.isArray(Domains) || Domains.length === 0 || Domains.includes(host)) {
              return next();
            }
          }
        }

        // 2) Bypass via header keys
        for (const { Name, Value, Domains } of checkpointConfig.BypassHeaderKeys) {
          // Get header value case-insensitively by checking all headers
          let headerVal = null;
          const headersMap = Object.fromEntries([...request.headers.entries()].map(([k, v]) => [k.toLowerCase(), v]));
          headerVal = headersMap[Name.toLowerCase()] || request.headers.get(Name);
          
          console.log(`DEBUG - Checking header ${Name}: received="${headerVal}", expected="${Value}", domains=${JSON.stringify(Domains)}`);
          
          if (headerVal === Value) {
            console.log(`DEBUG - Header value matched for ${Name}`);
            if (!Array.isArray(Domains) || Domains.length === 0 || Domains.includes(host)) {
              console.log(`DEBUG - Domain check passed for ${host}`);
              return next();
            } else {
              console.log(`DEBUG - Domain check failed: ${host} not in ${JSON.stringify(Domains)}`);
            }
          } else {
            console.log(`DEBUG - Header value mismatch for ${Name}`);
          }
        }

        // Handle token redirect for URL-token login
        const tokenResponse = await handleTokenRedirect(request);
        if (tokenResponse) {
          // Convert Response to Express response
          res.status(tokenResponse.status);
          tokenResponse.headers.forEach((value, key) => {
            res.setHeader(key, value);
          });
          const body = await tokenResponse.text();
          return res.send(body);
        }

        // Setup request context
        const url = new URL(request.url);
        let path = url.pathname;
        if (checkpointConfig.SanitizeURLs) {
          path = sanitizePath(path);
        }
        const method = request.method;

        // Always allow challenge & verify endpoints
        if (method === 'GET' && path === '/api/challenge') {
          const response = await handleGetCheckpointChallenge(request);
          res.status(response.status);
          response.headers.forEach((value, key) => {
            res.setHeader(key, value);
          });
          const body = await response.text();
          return res.send(body);
        }
        if (method === 'POST' && path === '/api/verify') {
          const response = await handleVerifyCheckpoint(request);
          res.status(response.status);
          response.headers.forEach((value, key) => {
            res.setHeader(key, value);
          });
          const body = await response.text();
          return res.send(body);
        }

        // Check new exclusion rules
        if (checkpointConfig.ExclusionRules && checkpointConfig.ExclusionRules.length > 0) {
          for (const rule of checkpointConfig.ExclusionRules) {
            // Check if path matches
            if (!rule.Path || !path.startsWith(rule.Path)) {
              continue;
            }

            // Check if host matches (if specified)
            if (rule.Hosts && rule.Hosts.length > 0 && !rule.Hosts.includes(host)) {
              continue;
            }

            // Check if user agent matches (if specified)
            if (rule.UserAgents && rule.UserAgents.length > 0) {
              const matchesUA = rule.UserAgents.some((ua) => userAgent.includes(ua));
              if (!matchesUA) {
                continue;
              }
            }

            // All conditions match - exclude this request
            return next();
          }
        }

        // Skip checkpoint for requests that don't accept HTML
        if (!req.accepts('html')) {
          return next();
        }

        // Validate session token
        const cookies = cookie.parse(request.headers.get('cookie') || '');
        const tokenCookie = cookies[checkpointConfig.CookieName];
        const validation = await validateToken(tokenCookie, request);
        if (validation) {
          // Active session: bypass checkpoint
          return next();
        }

        // Log new checkpoint flow
        console.log(`checkpoint: incoming ${method} ${request.url}`);
        console.log(`checkpoint: tokenCookie=${tokenCookie}`);
        console.log(`checkpoint: validateToken => ${validation}`);

        // Serve interstitial challenge
        const response = await serveInterstitial(request);
        res.status(response.status);
        response.headers.forEach((value, key) => {
          res.setHeader(key, value);
        });
        const body = await response.text();
        return res.send(body);
      }
    ]
  };
}

async function addToken(tokenHash, data) {
  if (!db) return;
  try {
    const ttlMs = checkpointConfig.TokenExpiration;

    await db.put(tokenHash, data);

    tokenExpirations.set(tokenHash, Date.now() + ttlMs);
  } catch (err) {
    console.error('Error adding token:', err);
  }
}

async function updateTokenVerification(tokenHash) {
  if (!db) return;
  try {
    const data = await db.get(tokenHash);
    data.LastVerified = new Date().toISOString();
    await db.put(tokenHash, data);
  } catch (err) {
    console.error('Error updating token verification:', err);
  }
}

async function lookupTokenData(tokenHash) {
  if (!db) return { data: null, found: false };
  try {
    const expireTime = tokenExpirations.get(tokenHash);
    if (!expireTime || expireTime <= Date.now()) {
      if (expireTime) {
        tokenExpirations.delete(tokenHash);
        try {
          await db.del(tokenHash);
        } catch (e) {}
      }
      return { data: null, found: false };
    }

    const data = await db.get(tokenHash);
    return { data, found: true };
  } catch (err) {
    if (err.code === 'LEVEL_NOT_FOUND') return { data: null, found: false };
    console.error('Error looking up token:', err);
    throw err;
  }
}

async function closeTokenStore() {
  if (db) await db.close();
}

function startCleanupTimer() {
  // Cleanup expired data hourly
  setInterval(() => {
    cleanupExpiredData();
  }, 3600000);
  // Cleanup expired challenges at the challenge expiration interval
  const challengeInterval = checkpointConfig.ChallengeExpiration || 60000;
  setInterval(() => {
    cleanupExpiredChallenges();
  }, challengeInterval);
}

function cleanupExpiredData() {
  const now = Date.now();
  let count = 0;

  try {
    for (const [nonce, ts] of usedNonces.entries()) {
      if (now - ts > checkpointConfig.MaxNonceAge) {
        usedNonces.delete(nonce);
        count++;
      }
    }
    if (count) console.log(`Checkpoint: cleaned up ${count} expired nonces.`);
  } catch (err) {
    console.error('Error cleaning up nonces:', err);
  }

  // Clean up expired tokens from cache
  let tokenCacheCount = 0;
  try {
    for (const [tokenKey, _] of tokenCache.entries()) {
      const expireTime = tokenExpirations.get(tokenKey);
      if (!expireTime || expireTime <= now) {
        tokenCache.delete(tokenKey);
        tokenExpirations.delete(tokenKey);
        tokenCacheCount++;
      }
    }
    if (tokenCacheCount)
      console.log(`Checkpoint: cleaned up ${tokenCacheCount} expired tokens from cache.`);
  } catch (err) {
    console.error('Error cleaning up token cache:', err);
  }

  try {
    ipRateLimit.clear();
    console.log('Checkpoint: IP rate limits reset.');
  } catch (err) {
    console.error('Error resetting IP rate limits:', err);
  }
}

function cleanupExpiredChallenges() {
  const now = Date.now();
  let count = 0;
  for (const [id, params] of challengeStore.entries()) {
    if (params.ExpiresAt && params.ExpiresAt < now) {
      // Record failure for expired challenges that were never completed
      safeRecordEvent('checkpoint.failure', {
        reason: 'challenge_expired',
        ip: params.ClientIP,
        challenge_id: id.substring(0, 8), // Include partial ID for debugging
        age_ms: now - params.CreatedAt, // How long the challenge existed
        expiry_ms: checkpointConfig.ChallengeExpiration, // Configured expiry time
      });
      challengeStore.delete(id);
      count++;
    }
  }
  if (count) console.log(`Checkpoint: cleaned up ${count} expired challenges.`);
}

async function initSecret() {
  try {
    if (!checkpointConfig.SecretConfigPath) {
      checkpointConfig.SecretConfigPath = join(rootDir, 'data', 'checkpoint_secret.json');
    }

    const secretPath = checkpointConfig.SecretConfigPath;
    const exists = fs.existsSync(secretPath);

    if (exists) {
      const loaded = loadSecretFromFile();
      if (loaded) {
        hmacSecret = loaded;
        console.log(`Loaded existing HMAC secret from ${secretPath}`);
        return;
      }
    }

    hmacSecret = crypto.randomBytes(32);
    fs.mkdirSync(path.dirname(secretPath), { recursive: true });

    const secretCfg = {
      hmac_secret: hmacSecret.toString('base64'),
      created_at: new Date().toISOString(),
      updated_at: new Date().toISOString(),
    };

    fs.writeFileSync(secretPath, JSON.stringify(secretCfg), { mode: 0o600 });
    console.log(`Created and saved new HMAC secret to ${secretPath}`);
  } catch (err) {
    console.error('Error initializing secret:', err);

    hmacSecret = crypto.randomBytes(32);
  }
}

function loadSecretFromFile() {
  try {
    const data = fs.readFileSync(checkpointConfig.SecretConfigPath, 'utf8');
    const cfg = JSON.parse(data);
    const buf = Buffer.from(cfg.hmac_secret, 'base64');
    if (buf.length < 16) return null;

    cfg.updated_at = new Date().toISOString();
    fs.writeFileSync(checkpointConfig.SecretConfigPath, JSON.stringify(cfg), { mode: 0o600 });
    return buf;
  } catch (e) {
    console.warn('Could not load HMAC secret from file:', e);
    return null;
  }
}

(async function initialize() {
  await initConfig();
  await initSecret();
  initTokenStore();
  startCleanupTimer();

  // Only register plugin if enabled
  if (checkpointConfig.Enabled !== false) {
    registerPlugin('checkpoint', CheckpointMiddleware());
  } else {
    console.log('Checkpoint plugin disabled via configuration');
  }
})();

export { checkpointConfig, addToken, updateTokenVerification, lookupTokenData, closeTokenStore };