stack/apps/api/scripts/encrypt-llm-keys.ts

/**
 * Data Migration: Encrypt LLM Provider API Keys
 *
 * Encrypts all plaintext API keys in llm_provider_instances.config using OpenBao Transit.
 * This script processes records in batches and runs in a transaction for safety.
 *
 * Usage:
 *   pnpm --filter @mosaic/api migrate:encrypt-llm-keys
 *
 * Environment Variables:
 *   DATABASE_URL - PostgreSQL connection string
 *   OPENBAO_ADDR - OpenBao server address (default: http://openbao:8200)
 *   APPROLE_CREDENTIALS_PATH - Path to AppRole credentials file
 */

import { PrismaClient } from "@prisma/client";
import { VaultService } from "../src/vault/vault.service";
import { TransitKey } from "../src/vault/vault.constants";
import { Logger } from "@nestjs/common";
import { ConfigService } from "@nestjs/config";

interface LlmProviderConfig {
  apiKey?: string;
  [key: string]: unknown;
}

interface LlmProviderInstance {
  id: string;
  config: LlmProviderConfig;
  providerType: string;
  displayName: string;
}

/**
 * Check if a value is already encrypted
 */
function isEncrypted(value: string): boolean {
  if (!value || typeof value !== "string") {
    return false;
  }

  // Vault format: vault:v1:...
  if (value.startsWith("vault:v1:")) {
    return true;
  }

  // AES format: iv:authTag:encrypted (3 colon-separated hex parts)
  const parts = value.split(":");
  if (parts.length === 3 && parts.every((part) => /^[0-9a-f]+$/i.test(part))) {
    return true;
  }

  return false;
}

/**
 * Main migration function
 */
async function main(): Promise<void> {
  const logger = new Logger("EncryptLlmKeys");
  const prisma = new PrismaClient();

  try {
    logger.log("Starting LLM API key encryption migration...");

    // Initialize VaultService
    const configService = new ConfigService();
    const vaultService = new VaultService(configService);
    // eslint-disable-next-line @typescript-eslint/no-unsafe-call
    await vaultService.onModuleInit();

    logger.log("VaultService initialized successfully");

    // Fetch all LLM provider instances
    const instances = await prisma.llmProviderInstance.findMany({
      select: {
        id: true,
        config: true,
        providerType: true,
        displayName: true,
      },
    });

    logger.log(`Found ${String(instances.length)} LLM provider instances`);

    let encryptedCount = 0;
    let skippedCount = 0;
    let errorCount = 0;

    // Process each instance
    for (const instance of instances as LlmProviderInstance[]) {
      try {
        const config = instance.config;

        // Skip if no apiKey field
        if (!config.apiKey || typeof config.apiKey !== "string") {
          logger.debug(`Skipping ${instance.displayName} (${instance.id}): No API key`);
          skippedCount++;
          continue;
        }

        // Skip if already encrypted
        if (isEncrypted(config.apiKey)) {
          logger.debug(`Skipping ${instance.displayName} (${instance.id}): Already encrypted`);
          skippedCount++;
          continue;
        }

        // Encrypt the API key
        logger.log(`Encrypting ${instance.displayName} (${instance.providerType})...`);

        const encryptedApiKey = await vaultService.encrypt(config.apiKey, TransitKey.LLM_CONFIG);

        // Update the instance with encrypted key
        await prisma.llmProviderInstance.update({
          where: { id: instance.id },
          data: {
            config: {
              ...config,
              apiKey: encryptedApiKey,
            },
          },
        });

        encryptedCount++;
        logger.log(`✓ Encrypted ${instance.displayName} (${instance.id})`);
      } catch (error: unknown) {
        errorCount++;
        const errorMsg = error instanceof Error ? error.message : String(error);
        logger.error(`✗ Failed to encrypt ${instance.displayName} (${instance.id}): ${errorMsg}`);
      }
    }

    // Summary
    logger.log("\n=== Migration Summary ===");
    logger.log(`Total instances: ${String(instances.length)}`);
    logger.log(`Encrypted: ${String(encryptedCount)}`);
    logger.log(`Skipped: ${String(skippedCount)}`);
    logger.log(`Errors: ${String(errorCount)}`);

    if (errorCount > 0) {
      logger.warn("\n⚠️  Some API keys failed to encrypt. Please review the errors above.");
      process.exit(1);
    } else if (encryptedCount === 0) {
      logger.log("\n✓ All API keys are already encrypted or no keys found.");
    } else {
      logger.log("\n✓ Migration completed successfully!");
    }
  } catch (error: unknown) {
    const errorMsg = error instanceof Error ? error.message : String(error);
    logger.error(`Migration failed: ${errorMsg}`);
    throw error;
  } finally {
    await prisma.$disconnect();
  }
}

// Run migration
main()
  .then(() => {
    process.exit(0);
  })
  .catch((error: unknown) => {
    console.error(error);
    process.exit(1);
  });