codeflash-internal/js/common/src/token-functions.ts

"use server"
import util from "util"
import { createHash, type Hash, randomBytes } from "crypto"
import { prisma } from "./prisma-client"
import { organizationMemberRepository } from "./repositories/organization-member-repository"

async function generateRandomAPIKey(): Promise<string> {
  const randomBytesAsync = util.promisify(randomBytes)
  const token = (await randomBytesAsync(48)).toString("base64url")
  return "cf-" + token
}
export const VS_CODE_KEY_NAME = "vsc-ext-3Yg1NRCS@6"
export async function hashApiKey(token: string) {
  // TODO: Consider stripping the cf- prefix from the token before hashing. This would reduce the chance of
  //  a collision, and would also make it harder to brute force the token.
  const hash: Hash = createHash("sha384")
  hash.update(token)
  return hash.digest("base64url")
}

export async function genAndStoreAPITokenHash(
  keyName: string,
  userId: string,
  organizationId?: string,
): Promise<string> {
  // The idea here is to not store the API Keys in the database in plaintext. Because if we do that, then
  // the keys are vulnerable to internal and external database breaches. Instead, we store the hash
  // of the key in the database. When a user wants to use the key, they send the plaintext key to the server,
  // which hashes it and compares it to the hash in the database. If they match, then the key is valid.
  // The reason to not use bcrypt is that bcrypt is designed to be slow, which is good for passwords, but
  // not good for API keys. Also with salting, one can't simply search for a key in the DB.
  // Instead, we use SHA-384, which is designed to be fast. Rainbow table attacks are mitigated by using a long
  // random key, for which rainbow tables become unfeasible and the attacker would have to brute force the key.
  // The reason to use SHA-384 instead of SHA-256 is that SHA-384 has a longer output, which makes it much harder
  // to brute force.

  const token: string = await generateRandomAPIKey()
  const hashedToken = await hashApiKey(token)
  await prisma.cf_api_keys.create({
    data: {
      key: hashedToken,
      user_id: userId,
      organization_id: organizationId,
      suffix: token.slice(-4),
      name: keyName,
    },
  })
  return token
}
export async function genAndStoreAPITokenHashForVSC(userId: string, organizationId?: string) {
  const token: string = await generateRandomAPIKey()
  const hashedToken = await hashApiKey(token)
  const data = await prisma.cf_api_keys.findFirst({
    where: {
      user_id: userId,
      name: VS_CODE_KEY_NAME,
    },
  })
  if (data) {
    await prisma.cf_api_keys.update({
      where: {
        id: data.id,
      },
      data: {
        key: hashedToken,
        suffix: token.slice(-4),
        organization_id: organizationId ?? null,
      },
    })
  } else {
    await prisma.cf_api_keys.create({
      data: {
        key: hashedToken,
        user_id: userId,
        organization_id: organizationId ?? null,
        suffix: token.slice(-4),
        name: VS_CODE_KEY_NAME,
      },
    })
  }

  return token
}
export interface ApiKeyInfo {
  user_id: string | null
  organization_id: string | null
}

export async function userForAPIKey(key: string): Promise<null | string> {
  // TODO: Add a rate limiter to prevent brute force attacks.
  const hashedToken = await hashApiKey(key)
  try {
    // Find the API key - only fetch what we need
    const apiKey = await prisma.cf_api_keys.findUnique({
      where: { key: hashedToken },
      select: {
        user_id: true,
      },
    })

    if (apiKey == null) {
      console.log(`[Auth] API key not found (suffix: ${key.slice(-4)})`)
      return null
    }

    // Update last_used timestamp
    await prisma.cf_api_keys.update({
      where: { key: hashedToken },
      data: { last_used: new Date() },
    })

    console.log(`[Auth] API key authenticated successfully`)
    return apiKey.user_id
  } catch (e) {
    console.error(`[Auth] Error in userForAPIKey:`, e)
    return null
  }
}

/**
 * Get full API key info including organization_id
 */
export async function getApiKeyInfo(key: string): Promise<ApiKeyInfo | null> {
  const hashedToken = await hashApiKey(key)
  try {
    const apiKey = await prisma.cf_api_keys.findUnique({
      where: { key: hashedToken },
      select: {
        user_id: true,
        organization_id: true,
      },
    })

    if (apiKey == null) {
      console.log(`[Auth] API key not found (suffix: ${key.slice(-4)})`)
      return null
    }

    // Update last_used timestamp
    await prisma.cf_api_keys.update({
      where: { key: hashedToken },
      data: { last_used: new Date() },
    })

    return {
      user_id: apiKey.user_id,
      organization_id: apiKey.organization_id,
    }
  } catch (e) {
    console.error(`[Auth] Error in getApiKeyInfo:`, e)
    return null
  }
}

/**
 * Get effective privacy mode for an API key.
 * Privacy mode is enabled if:
 * - The API key belongs to an organization with subscription AND org privacy_mode is enabled, OR
 * - The user is on a paid tier (not free) AND user privacy_mode is enabled
 *
 * When privacy mode is enabled, data is fetched from GitHub instead of being stored.
 * When privacy mode is disabled, data is stored as plaintext.
 */
export async function getEffectivePrivacyMode(
  userId: string,
  organizationId: string | null,
): Promise<boolean> {
  // If API key has an organization, check if org has subscription AND privacy_mode is enabled
  if (organizationId) {
    const org = await prisma.organizations.findUnique({
      where: { id: organizationId },
      select: {
        subscription: true,
        privacy_mode: true,
      },
    })

    // If org has subscription, use org's privacy_mode setting
    if (org?.subscription) {
      if (org.privacy_mode) {
        console.log(`[Privacy] Organization has subscription and privacy mode enabled`)
        return true
      } else {
        console.log(`[Privacy] Organization has subscription but privacy mode is disabled`)
        return false
      }
    }
  }

  // Check if user is on a paid tier AND has privacy mode enabled
  const [subscription, user] = await Promise.all([
    prisma.subscriptions.findUnique({
      where: { user_id: userId },
      select: { plan_type: true },
    }),
    prisma.users.findUnique({
      where: { user_id: userId },
      select: { privacy_mode: true },
    }),
  ])

  const isPaidUser = subscription?.plan_type && subscription.plan_type.toLowerCase() !== "free"

  if (isPaidUser && user?.privacy_mode) {
    return true
  }

  return false
}

export async function deleteAPIKeyById(id: number, userId: string) {
  return prisma.cf_api_keys.delete({
    where: {
      id,
      user_id: userId,
    },
  })
}

/**
 * Delete all API keys created by a user for a specific organization.
 * This should be called when a user is removed from an organization.
 */
export async function deleteOrganizationMemberApiKeys(
  userId: string,
  organizationId: string,
) {
  return prisma.cf_api_keys.deleteMany({
    where: {
      user_id: userId,
      organization_id: organizationId,
    },
  })
}

/**
 * Check whether a user has fewer than 30 API tokens.
 *
 */

export async function canInsertMoreTokens(userId: string): Promise<boolean> {
  const tokenCount = await prisma.cf_api_keys.count({
    where: {
      user_id: userId,
      name: { not: VS_CODE_KEY_NAME },
    },
  })
  return tokenCount < 30
}

/**
 * Wrapper function to enforce token limit consistently.
 * Call this instead of genAndStoreAPITokenHash if you want to prevent generating tokens
 * for users who already have 30 or more.
 */
export async function safeGenAndStoreAPITokenHash(
  keyName: string,
  userId: string,
  organizationId?: string,
): Promise<string> {
  if (!(await canInsertMoreTokens(userId))) {
    throw new Error("Token limit exceeded")
  }

  if (organizationId) {
    const isMember = await organizationMemberRepository.isMember(organizationId, userId)
    if (!isMember) {
      throw new Error("User is not a member of the specified organization")
    }
  }

  return genAndStoreAPITokenHash(keyName, userId, organizationId)
}