droid2api/auth.js

import fs from 'fs';
import path from 'path';
import os from 'os';
import { logDebug, logError, logInfo } from './logger.js';
import { getNextProxyAgent } from './proxy-manager.js';
import { getAccountManager, initializeAccountManager } from './account-manager.js';
import { getRefreshConfig, requestRefreshToken } from './refresh-client.js';

// State management for API key and refresh
let currentApiKey = null;
let currentRefreshToken = null;
let lastRefreshTime = null;
let tokenExpiresAt = null; // Token 过期时间戳 (ms)
let clientId = null;
let authSource = null; // 'env' or 'file' or 'factory_key' or 'client' or 'multi_account'
let authFilePath = null;
let factoryApiKey = null; // From FACTORY_API_KEY environment variable
let multiAccountMode = false; // 是否启用多账号模式
let lastSelectedAccountId = null; // 记录最后选择的账号ID，用于结果回调
let refreshInFlight = null; // 刷新锁，避免并发刷新

const REFRESH_URL = 'https://api.workos.com/user_management/authenticate';
const REFRESH_INTERVAL_HOURS = 6; // Refresh every 6 hours
const TOKEN_VALID_HOURS = 8; // Token valid for 8 hours
const REFRESH_BUFFER_MS = 30 * 60 * 1000; // 提前 30 分钟刷新

/**
 * Generate a ULID (Universally Unique Lexicographically Sortable Identifier)
 * Format: 26 characters using Crockford's Base32
 * First 10 chars: timestamp (48 bits)
 * Last 16 chars: random (80 bits)
 */
function generateULID() {
  // Crockford's Base32 alphabet (no I, L, O, U to avoid confusion)
  const ENCODING = '0123456789ABCDEFGHJKMNPQRSTVWXYZ';

  // Get timestamp in milliseconds
  const timestamp = Date.now();

  // Encode timestamp to 10 characters
  let time = '';
  let ts = timestamp;
  for (let i = 9; i >= 0; i--) {
    const mod = ts % 32;
    time = ENCODING[mod] + time;
    ts = Math.floor(ts / 32);
  }

  // Generate 16 random characters
  let randomPart = '';
  for (let i = 0; i < 16; i++) {
    const rand = Math.floor(Math.random() * 32);
    randomPart += ENCODING[rand];
  }

  return time + randomPart;
}

/**
 * Generate a client ID in format: client_01{ULID}
 */
function generateClientId() {
  const ulid = generateULID();
  return `client_01${ulid}`;
}

/**
 * Load auth configuration with priority system
 * Priority: accounts.json (multi-account) > FACTORY_API_KEY > refresh token mechanism > client authorization
 */
function loadAuthConfig() {
  // 0. Check accounts.json for multi-account mode (highest priority)
  const accountsPath = path.join(process.cwd(), 'accounts.json');
  logDebug(`Checking accounts.json at: ${accountsPath}`);
  try {
    if (fs.existsSync(accountsPath)) {
      logDebug('accounts.json exists, reading...');
      const accountsContent = fs.readFileSync(accountsPath, 'utf-8');
      const accountsData = JSON.parse(accountsContent);
      logDebug(`accounts.json parsed, accounts: ${accountsData.accounts?.length || 0}`);

      if (accountsData.accounts && accountsData.accounts.length > 0) {
        const activeAccounts = accountsData.accounts.filter(acc => acc.status === 'active');
        if (activeAccounts.length > 0) {
          logInfo(`Found accounts.json with ${activeAccounts.length} active account(s), enabling multi-account mode`);
          multiAccountMode = true;
          authSource = 'multi_account';
          return { type: 'multi_account', value: accountsPath };
        }
      }
    }
  } catch (error) {
    logError('Error reading accounts.json', error);
  }

  // 1. Check FACTORY_API_KEY environment variable
  const factoryKey = process.env.FACTORY_API_KEY;
  if (factoryKey && factoryKey.trim() !== '') {
    logInfo('Using fixed API key from FACTORY_API_KEY environment variable');
    factoryApiKey = factoryKey.trim();
    authSource = 'factory_key';
    return { type: 'factory_key', value: factoryKey.trim() };
  }

  // 2. Check refresh token mechanism (DROID_REFRESH_KEY)
  const envRefreshKey = process.env.DROID_REFRESH_KEY;
  if (envRefreshKey && envRefreshKey.trim() !== '') {
    logInfo('Using refresh token from DROID_REFRESH_KEY environment variable');
    authSource = 'env';
    authFilePath = path.join(process.cwd(), 'auth.json');
    return { type: 'refresh', value: envRefreshKey.trim() };
  }

  // 3. Check ~/.factory/auth.json
  const homeDir = os.homedir();
  const factoryAuthPath = path.join(homeDir, '.factory', 'auth.json');

  try {
    if (fs.existsSync(factoryAuthPath)) {
      const authContent = fs.readFileSync(factoryAuthPath, 'utf-8');
      const authData = JSON.parse(authContent);

      if (authData.refresh_token && authData.refresh_token.trim() !== '') {
        logInfo('Using refresh token from ~/.factory/auth.json');
        authSource = 'file';
        authFilePath = factoryAuthPath;

        // Also load access_token if available and not expired
        if (authData.access_token) {
          const expiresAt = authData.expires_at ? new Date(authData.expires_at).getTime() : null;
          const now = Date.now();

          if (expiresAt && expiresAt > now + REFRESH_BUFFER_MS) {
            // Token 还有效（且距离过期超过30分钟）
            currentApiKey = authData.access_token.trim();
            tokenExpiresAt = expiresAt;
            lastRefreshTime = authData.last_updated ? new Date(authData.last_updated).getTime() : now;
            logInfo(`Loaded valid token from file, expires at: ${new Date(expiresAt).toISOString()}`);
          } else if (expiresAt) {
            // Token 已过期或即将过期
            logInfo(`Stored token expired or expiring soon (expires_at: ${authData.expires_at}), will refresh`);
          } else {
            // 没有过期时间记录，按旧逻辑处理
            currentApiKey = authData.access_token.trim();
            logInfo('Loaded token from file (no expiry info, will check on first use)');
          }
        }

        return { type: 'refresh', value: authData.refresh_token.trim() };
      }
    }
  } catch (error) {
    logError('Error reading ~/.factory/auth.json', error);
  }

  // 4. No configured auth found - will use client authorization
  logInfo('No auth configuration found, will use client authorization headers');
  authSource = 'client';
  return { type: 'client', value: null };
}

/**
 * Refresh API key using refresh token
 */
async function refreshApiKey() {
  if (!currentRefreshToken) {
    throw new Error('No refresh token available');
  }

  if (refreshInFlight) {
    return refreshInFlight;
  }

  refreshInFlight = (async () => {
    if (!clientId) {
      clientId = 'client_01HNM792M5G5G1A2THWPXKFMXB';
      logDebug(`Using fixed client ID: ${clientId}`);
    }

    logInfo('Refreshing API key...');

    try {
      const proxyAgentInfo = getNextProxyAgent(REFRESH_URL);
      const refreshConfig = getRefreshConfig();
      const data = await requestRefreshToken({
        refreshUrl: REFRESH_URL,
        refreshToken: currentRefreshToken,
        clientId,
        proxyAgentInfo,
        ...refreshConfig
      });

      // Update tokens
      currentApiKey = data.access_token;
      currentRefreshToken = data.refresh_token;
      lastRefreshTime = Date.now();
      // 设置过期时间（默认8小时）
      tokenExpiresAt = lastRefreshTime + TOKEN_VALID_HOURS * 60 * 60 * 1000;

      // Log user info
      if (data.user) {
        logInfo(`Authenticated as: ${data.user.email} (${data.user.first_name} ${data.user.last_name})`);
        logInfo(`User ID: ${data.user.id}`);
        logInfo(`Organization ID: ${data.organization_id}`);
      }

      // Save tokens to file
      saveTokens(data.access_token, data.refresh_token);

      logInfo(`New Refresh-Key: ${currentRefreshToken}`);
      logInfo(`Token expires at: ${new Date(tokenExpiresAt).toISOString()}`);
      logInfo('API key refreshed successfully');
      return data.access_token;

    } catch (error) {
      logError('Failed to refresh API key', error);
      throw error;
    }
  })();

  try {
    return await refreshInFlight;
  } finally {
    refreshInFlight = null;
  }
}

/**
 * Save tokens to appropriate file
 * @param {string} accessToken - Access token to save
 * @param {string} refreshToken - Refresh token to save
 * @param {number} expiresInMs - Token validity duration in milliseconds (default: TOKEN_VALID_HOURS)
 */
function saveTokens(accessToken, refreshToken, expiresInMs = TOKEN_VALID_HOURS * 60 * 60 * 1000) {
  try {
    const now = Date.now();
    const expiresAt = new Date(now + expiresInMs).toISOString();

    const authData = {
      access_token: accessToken,
      refresh_token: refreshToken,
      expires_at: expiresAt,
      last_updated: new Date(now).toISOString()
    };

    // Ensure directory exists
    const dir = path.dirname(authFilePath);
    if (!fs.existsSync(dir)) {
      fs.mkdirSync(dir, { recursive: true });
    }

    // If saving to ~/.factory/auth.json, preserve other fields
    if (authSource === 'file' && fs.existsSync(authFilePath)) {
      try {
        const existingData = JSON.parse(fs.readFileSync(authFilePath, 'utf-8'));
        Object.assign(authData, existingData, {
          access_token: accessToken,
          refresh_token: refreshToken,
          expires_at: expiresAt,
          last_updated: authData.last_updated
        });
      } catch (error) {
        logError('Error reading existing auth file, will overwrite', error);
      }
    }

    fs.writeFileSync(authFilePath, JSON.stringify(authData, null, 2), 'utf-8');
    logDebug(`Tokens saved to ${authFilePath}`);

  } catch (error) {
    logError('Failed to save tokens', error);
  }
}

/**
 * Check if API key needs refresh
 * Uses actual expiration time if available, falls back to time-based check
 */
function shouldRefresh() {
  const now = Date.now();

  // 如果有过期时间，使用过期时间判断（提前30分钟刷新）
  if (tokenExpiresAt) {
    const shouldRefreshByExpiry = now + REFRESH_BUFFER_MS >= tokenExpiresAt;
    if (shouldRefreshByExpiry) {
      logDebug(`Token expiring soon (expires_at: ${new Date(tokenExpiresAt).toISOString()})`);
    }
    return shouldRefreshByExpiry;
  }

  // 回退到基于刷新时间的判断
  if (!lastRefreshTime) {
    return true;
  }

  const hoursSinceRefresh = (now - lastRefreshTime) / (1000 * 60 * 60);
  return hoursSinceRefresh >= REFRESH_INTERVAL_HOURS;
}

/**
 * Initialize auth system - load auth config and setup initial API key if needed
 */
export async function initializeAuth() {
  try {
    const authConfig = loadAuthConfig();

    if (authConfig.type === 'multi_account') {
      // Multi-account mode - initialize AccountManager
      const loaded = await initializeAccountManager();
      if (loaded) {
        const manager = getAccountManager();
        logInfo(`Auth system initialized with multi-account mode (${manager.getAccountCount()} accounts)`);
      } else {
        logError('Failed to initialize AccountManager, falling back to client authorization');
        authSource = 'client';
        multiAccountMode = false;
      }
    } else if (authConfig.type === 'factory_key') {
      // Using fixed FACTORY_API_KEY, no refresh needed
      logInfo('Auth system initialized with fixed API key');
    } else if (authConfig.type === 'refresh') {
      // Using refresh token mechanism
      currentRefreshToken = authConfig.value;

      // Try to refresh on startup to get fresh token
      try {
        await refreshApiKey();
        logInfo('Auth system initialized with refresh token mechanism');
      } catch (refreshError) {
        logError('Failed to refresh token on startup, falling back to client authorization', refreshError);
        authSource = 'client';
        logInfo('Auth system fallback to client authorization mode');
      }
    } else {
      // Using client authorization, no setup needed
      logInfo('Auth system initialized for client authorization mode');
    }

    logInfo('Auth system initialized successfully');
  } catch (error) {
    logError('Failed to initialize auth system', error);
    // Don't throw error, allow server to start with client authorization
    authSource = 'client';
    logInfo('Auth system fallback to client authorization mode');
  }
}

/**
 * Get API key based on configured authorization method
 * @param {string} clientAuthorization - Authorization header from client request (optional)
 */
export async function getApiKey(clientAuthorization = null) {
  // Priority 0: Multi-account mode
  if (authSource === 'multi_account' && multiAccountMode) {
    const manager = getAccountManager();
    const account = manager.selectAccount();
    lastSelectedAccountId = account.id;

    const accessToken = await manager.getAccessToken(account);
    return `Bearer ${accessToken}`;
  }

  // Priority 1: FACTORY_API_KEY environment variable
  if (authSource === 'factory_key' && factoryApiKey) {
    return `Bearer ${factoryApiKey}`;
  }

  // Priority 2: Refresh token mechanism
  if (authSource === 'env' || authSource === 'file') {
    // Check if we need to refresh
    if (shouldRefresh()) {
      logInfo('API key needs refresh (6+ hours old)');
      await refreshApiKey();
    }

    if (!currentApiKey) {
      throw new Error('No API key available from refresh token mechanism.');
    }

    return `Bearer ${currentApiKey}`;
  }

  // Priority 3: Client authorization header
  if (clientAuthorization) {
    logDebug('Using client authorization header');
    return clientAuthorization;
  }

  // No authorization available
  throw new Error('No authorization available. Please configure accounts.json, FACTORY_API_KEY, refresh token, or provide client authorization.');
}

/**
 * Record authentication result for multi-account mode
 * @param {string} endpoint - The API endpoint called
 * @param {boolean} success - Whether the request was successful
 * @param {number} statusCode - HTTP status code
 */
export function recordAuthResult(endpoint, success, statusCode) {
  if (authSource === 'multi_account' && multiAccountMode && lastSelectedAccountId) {
    const manager = getAccountManager();
    manager.recordResult(lastSelectedAccountId, endpoint, success, statusCode);
  }
}

/**
 * Get auth status for /status endpoint
 */
export function getAuthStatus() {
  if (authSource === 'multi_account' && multiAccountMode) {
    const manager = getAccountManager();
    return {
      mode: 'multi_account',
      ...manager.getStatus()
    };
  }

  return {
    mode: authSource,
    multiAccountMode: false
  };
}