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 os from "os";
import path from "path";
import { promisify } from "util";

import { TSshCertificateTemplates } from "@app/db/schemas";
import { BadRequestError } from "@app/lib/errors";
import { ms } from "@app/lib/ms";
import { CharacterType, characterValidator } from "@app/lib/validator/validate-string";
import { CertKeyAlgorithm } from "@app/services/certificate/certificate-types";

import {
  isValidHostPattern,
  isValidUserPattern
} from "../ssh-certificate-template/ssh-certificate-template-validators";
import { SshCertType, 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: CertKeyAlgorithm) => {
  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;

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

  try {
    // Generate the SSH key pair
    // The "-N ''" sets an empty passphrase
    // The keys are created in the temporary directory
    await execFileAsync("ssh-keygen", ["-t", keyType, "-b", keyBits, "-f", privateKeyFile, "-N", ""], {
      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]);
  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(() => {});
  }
};

/**
 * Create an SSH certificate for a user or host.
 */
export const createSshCert = async ({
  template,
  caPrivateKey,
  clientPublicKey,
  keyId,
  principals,
  requestedTtl,
  certType
}: TCreateSshCertDTO) => {
  // 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
  const ttl = validateSshCertificateTtl(template, requestedTtl);

  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(() => {});
  }
};