infisical/backend/src/ee/services/ssh/ssh-certificate-authority-fns.ts

import { execFile } from "child_process";
import crypto from "crypto";
import { promises as fs } from "fs";
import { Knex } from "knex";
import os from "os";
import path from "path";
import { promisify } from "util";

import { TSshCertificateTemplates } from "@app/db/schemas";
import { SshCertKeyAlgorithm } from "@app/ee/services/ssh-certificate/ssh-certificate-types";
import { BadRequestError } from "@app/lib/errors";
import { ms } from "@app/lib/ms";
import { CharacterType, characterValidator } from "@app/lib/validator/validate-string";
import { ActorType } from "@app/services/auth/auth-type";
import { KmsDataKey } from "@app/services/kms/kms-types";

import {
  isValidHostPattern,
  isValidUserPattern
} from "../ssh-certificate-template/ssh-certificate-template-validators";
import {
  SshCaKeySource,
  SshCaStatus,
  SshCertType,
  TConvertActorToPrincipalsDTO,
  TCreateSshCaHelperDTO,
  TCreateSshCertDTO
} from "./ssh-certificate-authority-types";

const execFileAsync = promisify(execFile);

const EXEC_TIMEOUT_MS = 10000; // 10 seconds
/* eslint-disable no-bitwise */
export const createSshCertSerialNumber = () => {
  const randomBytes = crypto.randomBytes(8); // 8 bytes = 64 bits
  randomBytes[0] &= 0x7f; // Ensure the most significant bit is 0 (to stay within unsigned range)
  return BigInt(`0x${randomBytes.toString("hex")}`).toString(10); // Convert to decimal
};

/**
 * Return a pair of SSH CA keys based on the specified key algorithm [keyAlgorithm].
 * We use this function because the key format generated by `ssh-keygen` is unique.
 */
export const createSshKeyPair = async (keyAlgorithm: SshCertKeyAlgorithm) => {
  const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "ssh-key-"));
  const privateKeyFile = path.join(tempDir, "id_key");
  const publicKeyFile = `${privateKeyFile}.pub`;

  let keyType: string;
  let keyBits: string | null;

  switch (keyAlgorithm) {
    case SshCertKeyAlgorithm.RSA_2048:
      keyType = "rsa";
      keyBits = "2048";
      break;
    case SshCertKeyAlgorithm.RSA_4096:
      keyType = "rsa";
      keyBits = "4096";
      break;
    case SshCertKeyAlgorithm.ECDSA_P256:
      keyType = "ecdsa";
      keyBits = "256";
      break;
    case SshCertKeyAlgorithm.ECDSA_P384:
      keyType = "ecdsa";
      keyBits = "384";
      break;
    case SshCertKeyAlgorithm.ED25519:
      keyType = "ed25519";
      keyBits = null;
      break;
    default:
      throw new BadRequestError({
        message: "Failed to produce SSH CA key pair generation command due to unrecognized key algorithm"
      });
  }

  try {
    const args = ["-t", keyType];
    if (keyBits !== null) {
      args.push("-b", keyBits);
    }
    args.push("-f", privateKeyFile, "-N", "");

    // Generate the SSH key pair
    // The "-N ''" sets an empty passphrase
    // The keys are created in the temporary directory
    await execFileAsync("ssh-keygen", args, {
      timeout: EXEC_TIMEOUT_MS
    });

    // Read the generated keys
    const publicKey = await fs.readFile(publicKeyFile, "utf8");
    const privateKey = await fs.readFile(privateKeyFile, "utf8");

    return { publicKey, privateKey };
  } finally {
    // Cleanup the temporary directory and all its contents
    await fs.rm(tempDir, { recursive: true, force: true }).catch(() => {});
  }
};

/**
 * Return the SSH public key for the given SSH private key.
 */
export const getSshPublicKey = async (privateKey: string) => {
  const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "ssh-key-"));
  const privateKeyFile = path.join(tempDir, "id_key");
  try {
    await fs.writeFile(privateKeyFile, privateKey, { mode: 0o600 });

    // Run ssh-keygen to extract the public key
    const { stdout } = await execFileAsync("ssh-keygen", ["-y", "-f", privateKeyFile], {
      encoding: "utf8",
      timeout: EXEC_TIMEOUT_MS
    });
    return stdout.trim();
  } finally {
    // Ensure that files and the temporary directory are cleaned up
    await fs.rm(tempDir, { recursive: true, force: true }).catch(() => {});
  }
};

/**
 * Validate the requested SSH certificate type based on the SSH certificate template configuration.
 */
export const validateSshCertificateType = (template: TSshCertificateTemplates, certType: SshCertType) => {
  if (!template.allowUserCertificates && certType === SshCertType.USER) {
    throw new BadRequestError({ message: "Failed to validate user certificate type due to template restriction" });
  }

  if (!template.allowHostCertificates && certType === SshCertType.HOST) {
    throw new BadRequestError({ message: "Failed to validate host certificate type due to template restriction" });
  }
};

/**
 * Validate the requested SSH certificate principals based on the SSH certificate template configuration.
 */
export const validateSshCertificatePrincipals = (
  certType: SshCertType,
  template: TSshCertificateTemplates,
  principals: string[]
) => {
  /**
   * Validate and sanitize a principal string
   */
  const validatePrincipal = (principal: string) => {
    const sanitized = principal.trim();

    // basic checks for empty or control characters
    if (sanitized.length === 0) {
      throw new BadRequestError({
        message: "Principal cannot be an empty string."
      });
    }

    if (/\r|\n|\t|\0/.test(sanitized)) {
      throw new BadRequestError({
        message: `Principal '${sanitized}' contains invalid whitespace or control characters.`
      });
    }

    // disallow whitespace anywhere
    if (/\s/.test(sanitized)) {
      throw new BadRequestError({
        message: `Principal '${sanitized}' cannot contain whitespace.`
      });
    }

    // restrict allowed characters to letters, digits, dot, underscore, and hyphen
    if (
      !characterValidator([
        CharacterType.AlphaNumeric,
        CharacterType.Period,
        CharacterType.Underscore,
        CharacterType.Hyphen
      ])(sanitized)
    ) {
      throw new BadRequestError({
        message: `Principal '${sanitized}' contains invalid characters. Allowed: alphanumeric, '.', '_', '-'.`
      });
    }

    // disallow leading hyphen to avoid potential argument-like inputs
    if (sanitized.startsWith("-")) {
      throw new BadRequestError({
        message: `Principal '${sanitized}' cannot start with a hyphen.`
      });
    }

    // length restriction (adjust as needed)
    if (sanitized.length > 64) {
      throw new BadRequestError({
        message: `Principal '${sanitized}' is too long.`
      });
    }

    return sanitized;
  };

  // Sanitize and validate all principals using the helper
  const sanitizedPrincipals = principals.map(validatePrincipal);

  switch (certType) {
    case SshCertType.USER: {
      if (template.allowedUsers.length === 0) {
        throw new BadRequestError({
          message: "No allowed users are configured in the SSH certificate template."
        });
      }

      const allowsAllUsers = template.allowedUsers.includes("*") ?? false;

      sanitizedPrincipals.forEach((principal) => {
        if (principal === "*") {
          throw new BadRequestError({
            message: `Principal '*' is not allowed for user certificates.`
          });
        }
        if (allowsAllUsers && !isValidUserPattern(principal)) {
          throw new BadRequestError({
            message: `Principal '${principal}' does not match a valid user pattern.`
          });
        }
        if (!allowsAllUsers && !template.allowedUsers.includes(principal)) {
          throw new BadRequestError({
            message: `Principal '${principal}' is not in the list of allowed users.`
          });
        }
      });
      break;
    }
    case SshCertType.HOST: {
      if (template.allowedHosts.length === 0) {
        throw new BadRequestError({
          message: "No allowed hosts are configured in the SSH certificate template."
        });
      }

      const allowsAllHosts = template.allowedHosts.includes("*") ?? false;

      sanitizedPrincipals.forEach((principal) => {
        if (principal.includes("*")) {
          throw new BadRequestError({
            message: `Principal '${principal}' with wildcards is not allowed for host certificates.`
          });
        }
        if (allowsAllHosts && !isValidHostPattern(principal)) {
          throw new BadRequestError({
            message: `Principal '${principal}' does not match a valid host pattern.`
          });
        }

        if (
          !allowsAllHosts &&
          !template.allowedHosts.some((allowedHost) => {
            if (allowedHost.startsWith("*.")) {
              const baseDomain = allowedHost.slice(2); // Remove the leading "*."
              return principal.endsWith(`.${baseDomain}`);
            }
            return principal === allowedHost;
          })
        ) {
          throw new BadRequestError({
            message: `Principal '${principal}' is not in the list of allowed hosts or domains.`
          });
        }
      });
      break;
    }
    default:
      throw new BadRequestError({
        message: "Failed to validate SSH certificate principals due to unrecognized requested certificate type"
      });
  }
};

/**
 * Validate the requested SSH certificate TTL based on the SSH certificate template configuration.
 */
export const validateSshCertificateTtl = (template: TSshCertificateTemplates, ttl?: string) => {
  if (!ttl) {
    // use default template ttl
    return Math.ceil(ms(template.ttl) / 1000);
  }

  if (ms(ttl) > ms(template.maxTTL)) {
    throw new BadRequestError({
      message: "Failed TTL validation due to TTL being greater than configured max TTL on template"
    });
  }

  return Math.ceil(ms(ttl) / 1000);
};

/**
 * Validate the requested SSH certificate key ID to ensure
 * that it only contains alphanumeric characters with no spaces.
 */
export const validateSshCertificateKeyId = (keyId: string) => {
  const regex = characterValidator([
    CharacterType.AlphaNumeric,
    CharacterType.Hyphen,
    CharacterType.Colon,
    CharacterType.Period
  ]);
  if (!regex(keyId)) {
    throw new BadRequestError({
      message:
        "Failed to validate Key ID because it can only contain alphanumeric characters and hyphens, with no spaces."
    });
  }

  if (keyId.length > 50) {
    throw new BadRequestError({
      message: "keyId can only be up to 50 characters long."
    });
  }
};

/**
 * Validate the format of the SSH public key
 */
const validateSshPublicKey = async (publicKey: string) => {
  const validPrefixes = ["ssh-rsa", "ssh-ed25519", "ecdsa-sha2-nistp256", "ecdsa-sha2-nistp384"];
  const startsWithValidPrefix = validPrefixes.some((prefix) => publicKey.startsWith(`${prefix} `));
  if (!startsWithValidPrefix) {
    throw new BadRequestError({ message: "Failed to validate SSH public key format: unsupported key type." });
  }

  // write the key to a temp file and run `ssh-keygen -l -f`
  // check to see if OpenSSH can read/interpret the public key
  const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "ssh-pubkey-"));
  const pubKeyFile = path.join(tempDir, "key.pub");

  try {
    await fs.writeFile(pubKeyFile, publicKey, { mode: 0o600 });
    await execFileAsync("ssh-keygen", ["-l", "-f", pubKeyFile], { timeout: EXEC_TIMEOUT_MS });
  } catch (error) {
    throw new BadRequestError({
      message: "Failed to validate SSH public key format: could not be parsed."
    });
  } finally {
    await fs.rm(tempDir, { recursive: true, force: true }).catch(() => {});
  }
};

export const getKeyAlgorithmFromFingerprintOutput = (output: string): SshCertKeyAlgorithm | undefined => {
  const parts = output.trim().split(" ");
  const bitsInt = parseInt(parts[0], 10);
  const keyTypeRaw = parts.at(-1)?.replace(/[()]/g, ""); // remove surrounding parentheses

  if (keyTypeRaw === "RSA") {
    return bitsInt === 2048 ? SshCertKeyAlgorithm.RSA_2048 : SshCertKeyAlgorithm.RSA_4096;
  }

  if (keyTypeRaw === "ECDSA") {
    return bitsInt === 256 ? SshCertKeyAlgorithm.ECDSA_P256 : SshCertKeyAlgorithm.ECDSA_P384;
  }

  if (keyTypeRaw === "ED25519") {
    // TODO: test
    return SshCertKeyAlgorithm.ED25519;
  }

  return undefined;
};

export const normalizeSshPrivateKey = (raw: string): string => {
  return `${raw
    .replace(/\r\n/g, "\n") // Windows CRLF → LF
    .replace(/\r/g, "\n") // Old Mac CR → LF
    .replace(/\\n/g, "\n") // Double-escaped \n
    .trim()}\n`;
};

/**
 * Validate the format of the SSH private key
 *
 * Returns the SSH public key corresponding to the private key
 * and the key algorithm categorization.
 */
export const validateSshPrivateKey = async (privateKey: string) => {
  const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "ssh-privkey-"));
  const privateKeyFile = path.join(tempDir, "id_key");

  try {
    await fs.writeFile(privateKeyFile, privateKey, {
      encoding: "utf8",
      mode: 0o600
    });

    // This will fail if the private key is malformed or unreadable
    const { stdout: publicKey } = await execFileAsync("ssh-keygen", ["-y", "-f", privateKeyFile], {
      timeout: EXEC_TIMEOUT_MS
    });

    const { stdout: fingerprint } = await execFileAsync("ssh-keygen", ["-lf", privateKeyFile]);
    const keyAlgorithm = getKeyAlgorithmFromFingerprintOutput(fingerprint);

    if (!keyAlgorithm) {
      throw new BadRequestError({
        message: "Failed to validate SSH private key format: The key algorithm is not supported."
      });
    }

    return {
      publicKey,
      keyAlgorithm
    };
  } catch (err) {
    throw new BadRequestError({
      message: "Failed to validate SSH private key format: could not be parsed."
    });
  } finally {
    await fs.rm(tempDir, { recursive: true, force: true }).catch(() => {});
  }
};

/**
 * Validate that the provided public and private keys are valid and constitute
 * a matching SSH key pair.
 */
export const validateExternalSshCaKeyPair = async (publicKey: string, privateKey: string) => {
  await validateSshPublicKey(publicKey);

  const { publicKey: derivedPublicKey, keyAlgorithm } = await validateSshPrivateKey(privateKey);

  if (publicKey.trim() !== derivedPublicKey.trim()) {
    throw new BadRequestError({
      message: "Failed to validate matching SSH key pair."
    });
  }

  return keyAlgorithm;
};

/**
 * Create an SSH certificate for a user or host.
 */
export const createSshCert = async ({
  template,
  caPrivateKey,
  clientPublicKey,
  keyId,
  principals,
  requestedTtl, // in ms lib format
  certType
}: TCreateSshCertDTO) => {
  let ttl: number | undefined;

  if (!template && requestedTtl) {
    const parsedTtl = Math.ceil(ms(requestedTtl) / 1000);
    if (parsedTtl > 0) ttl = parsedTtl;
  }

  if (template) {
    // validate if the requested [certType] is allowed under the template configuration
    validateSshCertificateType(template, certType);

    // validate if the requested [principals] are valid for the given [certType] under the template configuration
    validateSshCertificatePrincipals(certType, template, principals);

    // validate if the requested TTL is valid under the template configuration
    ttl = validateSshCertificateTtl(template, requestedTtl);
  }

  if (!ttl) {
    throw new BadRequestError({
      message: "Failed to create SSH certificate due to missing TTL"
    });
  }

  validateSshCertificateKeyId(keyId);
  await validateSshPublicKey(clientPublicKey);

  const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "ssh-cert-"));

  const publicKeyFile = path.join(tempDir, "user_key.pub");
  const privateKeyFile = path.join(tempDir, "ca_key");
  const signedPublicKeyFile = path.join(tempDir, "user_key-cert.pub");

  const serialNumber = createSshCertSerialNumber();

  // Build `ssh-keygen` arguments for signing
  // Using an array avoids shell injection issues
  const sshKeygenArgs = [
    certType === "host" ? "-h" : null, // host certificate if needed
    "-s",
    privateKeyFile, // path to SSH CA private key
    "-I",
    keyId, // identity (key ID)
    "-n",
    principals.join(","), // principals
    "-V",
    `+${ttl}s`, // validity (TTL in seconds)
    "-z",
    serialNumber, // serial number
    publicKeyFile // public key file to sign
  ].filter(Boolean) as string[];

  try {
    // Write public and private keys to the temp directory
    await fs.writeFile(publicKeyFile, clientPublicKey, { mode: 0o600 });
    await fs.writeFile(privateKeyFile, caPrivateKey, { mode: 0o600 });

    // Execute the signing process
    await execFileAsync("ssh-keygen", sshKeygenArgs, { encoding: "utf8", timeout: EXEC_TIMEOUT_MS });

    // Read the signed public key from the generated cert file
    const signedPublicKey = await fs.readFile(signedPublicKeyFile, "utf8");

    return { serialNumber, signedPublicKey, ttl };
  } finally {
    // Cleanup the temporary directory and all its contents
    await fs.rm(tempDir, { recursive: true, force: true }).catch(() => {});
  }
};

export const createSshCaHelper = async ({
  projectId,
  friendlyName,
  keyAlgorithm: requestedKeyAlgorithm,
  keySource,
  externalPk,
  externalSk,
  sshCertificateAuthorityDAL,
  sshCertificateAuthoritySecretDAL,
  kmsService,
  tx: outerTx
}: TCreateSshCaHelperDTO) => {
  // Function to handle the actual creation logic
  const processCreation = async (tx: Knex) => {
    let publicKey: string;
    let privateKey: string;
    let keyAlgorithm: SshCertKeyAlgorithm = requestedKeyAlgorithm;
    if (keySource === SshCaKeySource.INTERNAL) {
      // generate SSH CA key pair internally
      ({ publicKey, privateKey } = await createSshKeyPair(requestedKeyAlgorithm));
    } else {
      // use external SSH CA key pair
      if (!externalPk || !externalSk) {
        throw new BadRequestError({
          message: "Public and private keys are required if generateSigningKey is false"
        });
      }
      publicKey = externalPk;
      privateKey = externalSk;
      keyAlgorithm = await validateExternalSshCaKeyPair(publicKey, privateKey);
    }
    const ca = await sshCertificateAuthorityDAL.create(
      {
        projectId,
        friendlyName,
        status: SshCaStatus.ACTIVE,
        keyAlgorithm,
        keySource
      },
      tx
    );
    const { encryptor: secretManagerEncryptor } = await kmsService.createCipherPairWithDataKey(
      {
        type: KmsDataKey.SecretManager,
        projectId
      },
      tx
    );
    await sshCertificateAuthoritySecretDAL.create(
      {
        sshCaId: ca.id,
        encryptedPrivateKey: secretManagerEncryptor({ plainText: Buffer.from(privateKey, "utf8") }).cipherTextBlob
      },
      tx
    );
    return { ...ca, publicKey };
  };

  if (outerTx) {
    return processCreation(outerTx);
  }

  return sshCertificateAuthorityDAL.transaction(processCreation);
};

/**
 * Convert an actor to a list of principals to be included in an SSH certificate.
 *
 * (dangtony98): This function is only supported for user actors at the moment and returns
 * only the email of the associated user. In the future, we will consider other
 * actor types and attributes such as group membership slugs and/or metadata to be
 * included in the list of principals.
 */
export const convertActorToPrincipals = async ({ userDAL, actor, actorId }: TConvertActorToPrincipalsDTO) => {
  if (actor !== ActorType.USER) {
    throw new BadRequestError({
      message: "Failed to convert actor to principals due to unsupported actor type"
    });
  }

  const user = await userDAL.findById(actorId);

  return [user.username];
};