Termix/src/backend/utils/ssh-key-utils.ts

import { createHash, createPrivateKey } from "node:crypto";
import sshpk from "sshpk";
import { getErrorMessage } from "./error-message.js";
import ssh2Pkg from "ssh2";
const ssh2Utils = ssh2Pkg.utils;
const { parsePrivateKey } = sshpk;

function detectKeyTypeFromContent(keyContent: string): string {
  const content = keyContent.trim();

  if (content.includes("-----BEGIN OPENSSH PRIVATE KEY-----")) {
    if (
      content.includes("ssh-ed25519") ||
      content.includes("AAAAC3NzaC1lZDI1NTE5")
    ) {
      return "ssh-ed25519";
    }
    if (content.includes("ssh-rsa") || content.includes("AAAAB3NzaC1yc2E")) {
      return "ssh-rsa";
    }
    if (content.includes("ecdsa-sha2-nistp256")) {
      return "ecdsa-sha2-nistp256";
    }
    if (content.includes("ecdsa-sha2-nistp384")) {
      return "ecdsa-sha2-nistp384";
    }
    if (content.includes("ecdsa-sha2-nistp521")) {
      return "ecdsa-sha2-nistp521";
    }

    try {
      const base64Content = content
        .replace("-----BEGIN OPENSSH PRIVATE KEY-----", "")
        .replace("-----END OPENSSH PRIVATE KEY-----", "")
        .replace(/\s/g, "");

      const decoded = Buffer.from(base64Content, "base64").toString("binary");

      if (decoded.includes("ssh-rsa")) {
        return "ssh-rsa";
      }
      if (decoded.includes("ssh-ed25519")) {
        return "ssh-ed25519";
      }
      if (decoded.includes("ecdsa-sha2-nistp256")) {
        return "ecdsa-sha2-nistp256";
      }
      if (decoded.includes("ecdsa-sha2-nistp384")) {
        return "ecdsa-sha2-nistp384";
      }
      if (decoded.includes("ecdsa-sha2-nistp521")) {
        return "ecdsa-sha2-nistp521";
      }

      return "ssh-rsa";
    } catch {
      return "ssh-rsa";
    }
  }

  if (content.includes("-----BEGIN RSA PRIVATE KEY-----")) {
    return "ssh-rsa";
  }
  if (content.includes("-----BEGIN DSA PRIVATE KEY-----")) {
    return "ssh-dss";
  }
  if (content.includes("-----BEGIN EC PRIVATE KEY-----")) {
    return "ecdsa-sha2-nistp256";
  }

  if (content.includes("-----BEGIN PRIVATE KEY-----")) {
    try {
      const base64Content = content
        .replace("-----BEGIN PRIVATE KEY-----", "")
        .replace("-----END PRIVATE KEY-----", "")
        .replace(/\s/g, "");

      const decoded = Buffer.from(base64Content, "base64");
      const decodedString = decoded.toString("binary");

      if (decodedString.includes("1.2.840.113549.1.1.1")) {
        return "ssh-rsa";
      } else if (decodedString.includes("1.2.840.10045.2.1")) {
        if (decodedString.includes("1.2.840.10045.3.1.7")) {
          return "ecdsa-sha2-nistp256";
        }
        return "ecdsa-sha2-nistp256";
      } else if (decodedString.includes("1.3.101.112")) {
        return "ssh-ed25519";
      }
    } catch {
      // expected - base64 decode may fail for some key formats
    }

    if (content.length < 800) {
      return "ssh-ed25519";
    } else if (content.length > 1600) {
      return "ssh-rsa";
    } else {
      return "ecdsa-sha2-nistp256";
    }
  }

  return "unknown";
}

function detectPublicKeyTypeFromContent(publicKeyContent: string): string {
  const content = publicKeyContent.trim();

  // OpenSSH certificate types (must be checked before plain key types)
  if (content.startsWith("ssh-ed25519-cert-v01@openssh.com ")) {
    return "ssh-ed25519-cert-v01@openssh.com";
  }
  if (content.startsWith("ssh-rsa-cert-v01@openssh.com ")) {
    return "ssh-rsa-cert-v01@openssh.com";
  }
  if (content.startsWith("ecdsa-sha2-nistp256-cert-v01@openssh.com ")) {
    return "ecdsa-sha2-nistp256-cert-v01@openssh.com";
  }
  if (content.startsWith("ecdsa-sha2-nistp384-cert-v01@openssh.com ")) {
    return "ecdsa-sha2-nistp384-cert-v01@openssh.com";
  }
  if (content.startsWith("ecdsa-sha2-nistp521-cert-v01@openssh.com ")) {
    return "ecdsa-sha2-nistp521-cert-v01@openssh.com";
  }
  if (content.startsWith("sk-ssh-ed25519-cert-v01@openssh.com ")) {
    return "sk-ssh-ed25519-cert-v01@openssh.com";
  }

  // Plain public keys
  if (content.startsWith("ssh-rsa ")) {
    return "ssh-rsa";
  }
  if (content.startsWith("ssh-ed25519 ")) {
    return "ssh-ed25519";
  }
  if (content.startsWith("ecdsa-sha2-nistp256 ")) {
    return "ecdsa-sha2-nistp256";
  }
  if (content.startsWith("ecdsa-sha2-nistp384 ")) {
    return "ecdsa-sha2-nistp384";
  }
  if (content.startsWith("ecdsa-sha2-nistp521 ")) {
    return "ecdsa-sha2-nistp521";
  }
  if (content.startsWith("ssh-dss ")) {
    return "ssh-dss";
  }

  if (content.includes("-----BEGIN PUBLIC KEY-----")) {
    try {
      const base64Content = content
        .replace("-----BEGIN PUBLIC KEY-----", "")
        .replace("-----END PUBLIC KEY-----", "")
        .replace(/\s/g, "");

      const decoded = Buffer.from(base64Content, "base64");
      const decodedString = decoded.toString("binary");

      if (decodedString.includes("1.2.840.113549.1.1.1")) {
        return "ssh-rsa";
      } else if (decodedString.includes("1.2.840.10045.2.1")) {
        if (decodedString.includes("1.2.840.10045.3.1.7")) {
          return "ecdsa-sha2-nistp256";
        }
        return "ecdsa-sha2-nistp256";
      } else if (decodedString.includes("1.3.101.112")) {
        return "ssh-ed25519";
      }
    } catch {
      // expected - base64 decode may fail for some key formats
    }

    if (content.length < 400) {
      return "ssh-ed25519";
    } else if (content.length > 600) {
      return "ssh-rsa";
    } else {
      return "ecdsa-sha2-nistp256";
    }
  }

  if (content.includes("-----BEGIN RSA PUBLIC KEY-----")) {
    return "ssh-rsa";
  }

  if (content.includes("AAAAB3NzaC1yc2E")) {
    return "ssh-rsa";
  }
  if (content.includes("AAAAC3NzaC1lZDI1NTE5")) {
    return "ssh-ed25519";
  }
  if (content.includes("AAAAE2VjZHNhLXNoYTItbmlzdHAyNTY")) {
    return "ecdsa-sha2-nistp256";
  }
  if (content.includes("AAAAE2VjZHNhLXNoYTItbmlzdHAzODQ")) {
    return "ecdsa-sha2-nistp384";
  }
  if (content.includes("AAAAE2VjZHNhLXNoYTItbmlzdHA1MjE")) {
    return "ecdsa-sha2-nistp521";
  }
  if (content.includes("AAAAB3NzaC1kc3M")) {
    return "ssh-dss";
  }

  return "unknown";
}

export interface KeyInfo {
  privateKey: string;
  publicKey: string;
  keyType: string;
  success: boolean;
  error?: string;
}

export interface PublicKeyInfo {
  publicKey: string;
  keyType: string;
  success: boolean;
  error?: string;
}

export interface KeyPairValidationResult {
  isValid: boolean;
  privateKeyType: string;
  publicKeyType: string;
  generatedPublicKey?: string;
  error?: string;
}

const PUTTY_PRIVATE_KEY_RE = /^PuTTY-User-Key-File-(\d+):\s*(.+)$/m;

function normalizePrivateKeyText(privateKeyData: string): string {
  return privateKeyData.trim().replace(/\r\n/g, "\n").replace(/\r/g, "\n");
}

function getUnsupportedPrivateKeyError(privateKeyData: string): string {
  const puttyMatch = PUTTY_PRIVATE_KEY_RE.exec(privateKeyData.trim());
  if (puttyMatch) {
    const [, version, type] = puttyMatch;
    return `Unsupported PuTTY PPK v${version} private key format (${type}). Convert the key to OpenSSH format with PuTTYgen, or use a PuTTY PPK v2 RSA/DSA key.`;
  }

  return "Unsupported private key format. Use an OpenSSH, PEM, or PuTTY PPK v2 RSA/DSA private key.";
}

const SECURITY_KEY_ERROR =
  "This is a security key (FIDO2) key file. It only holds a handle, so the key itself must sign on your device. Load it with ssh-add and use SSH agent authentication in the desktop app.";

/** An OpenSSH private key file for a FIDO2 security key (sk-ssh-ed25519, sk-ecdsa). */
export function isSecurityKeyFile(privateKeyData: string): boolean {
  const match = privateKeyData.match(
    /-----BEGIN OPENSSH PRIVATE KEY-----([\s\S]*?)-----END OPENSSH PRIVATE KEY-----/,
  );
  if (!match) return false;
  try {
    const body = Buffer.from(match[1].replace(/\s+/g, ""), "base64");
    return /sk-(?:ssh-ed25519|ecdsa-sha2-nistp256)@openssh\.com/.test(
      body.toString("latin1"),
    );
  } catch {
    return false;
  }
}

// Unlocking an OpenSSH key runs bcrypt on the main thread (about a second per
// try), and every connect used to do it 3 or more times. Keep the unlocked copy
// in memory for a short while so a burst of connects pays for it once.
const UNLOCKED_KEY_TTL_MS = 15 * 60 * 1000;
const UNLOCKED_KEY_MAX = 64;
const unlockedKeys = new Map<string, { text: string; expires: number }>();

function unlockOpenSSHKey(cleanKey: string, passphrase?: string): string {
  if (
    !passphrase ||
    !cleanKey.startsWith("-----BEGIN OPENSSH PRIVATE KEY-----")
  )
    return cleanKey;
  const id = createHash("sha256")
    .update(cleanKey)
    .update("\n")
    .update(passphrase)
    .digest("hex");
  const now = Date.now();
  const hit = unlockedKeys.get(id);
  if (hit && hit.expires > now) return hit.text;
  unlockedKeys.delete(id);
  let text: string;
  try {
    text = parsePrivateKey(cleanKey, "openssh", { passphrase }).toString(
      "openssh",
    );
  } catch {
    // Wrong passphrase or a cipher sshpk lacks; the normal path reports it.
    return cleanKey;
  }
  if (ssh2Utils.parseKey(text) instanceof Error) return cleanKey;
  if (unlockedKeys.size >= UNLOCKED_KEY_MAX) {
    const oldest = unlockedKeys.keys().next().value;
    if (oldest !== undefined) unlockedKeys.delete(oldest);
  }
  unlockedKeys.set(id, { text, expires: now + UNLOCKED_KEY_TTL_MS });
  return text;
}

export function clearUnlockedKeyCache(): void {
  unlockedKeys.clear();
}

export function preparePrivateKeyForSSH2(
  privateKeyData: string,
  passphrase?: string,
): Buffer {
  if (isSecurityKeyFile(privateKeyData)) throw new Error(SECURITY_KEY_ERROR);
  let cleanKey = unlockOpenSSHKey(
    normalizePrivateKeyText(privateKeyData),
    passphrase,
  );
  if (
    /^-----BEGIN (?:ENCRYPTED )?PRIVATE KEY-----/.test(cleanKey) &&
    ssh2Utils.parseKey(cleanKey, passphrase) instanceof Error
  ) {
    try {
      const pem = createPrivateKey({ key: cleanKey, format: "pem", passphrase })
        .export({ type: "pkcs8", format: "pem" })
        .toString();
      cleanKey = parsePrivateKey(pem, "pkcs8").toString("openssh");
    } catch {
      throw new Error(
        cleanKey.startsWith("-----BEGIN ENCRYPTED PRIVATE KEY-----")
          ? "Could not decrypt PKCS#8 private key; check the passphrase"
          : "Invalid or unsupported PKCS#8 private key",
      );
    }
    const parsed = ssh2Utils.parseKey(cleanKey);
    if (parsed instanceof Error) throw parsed;
  }
  const keyInfo = parseSSHKey(cleanKey, passphrase);

  if (!keyInfo.success) {
    throw new Error(keyInfo.error || getUnsupportedPrivateKeyError(cleanKey));
  }

  return Buffer.from(cleanKey, "utf8");
}

export function isPrivateKeyPassphraseError(error: unknown): boolean {
  return /passphrase/i.test(getErrorMessage(error, ""));
}

export function parseSSHKey(
  privateKeyData: string,
  passphrase?: string,
): KeyInfo {
  try {
    const cleanKey = normalizePrivateKeyText(privateKeyData);
    let keyType = "unknown";
    let publicKey = "";
    let useSSH2 = false;

    if (ssh2Utils && typeof ssh2Utils.parseKey === "function") {
      try {
        const parsedKey = ssh2Utils.parseKey(cleanKey, passphrase);

        if (parsedKey instanceof Error) {
          throw parsedKey;
        } else {
          if (parsedKey.type) {
            keyType = parsedKey.type;
          }

          try {
            const publicKeyBuffer = parsedKey.getPublicSSH();

            if (Buffer.isBuffer(publicKeyBuffer)) {
              const base64Data = publicKeyBuffer.toString("base64");

              if (keyType === "ssh-rsa") {
                publicKey = `ssh-rsa ${base64Data}`;
              } else if (keyType === "ssh-ed25519") {
                publicKey = `ssh-ed25519 ${base64Data}`;
              } else if (keyType.startsWith("ecdsa-")) {
                publicKey = `${keyType} ${base64Data}`;
              } else {
                publicKey = `${keyType} ${base64Data}`;
              }
            } else {
              publicKey = "";
            }
          } catch {
            publicKey = "";
          }

          useSSH2 = true;
        }
      } catch {
        // expected - ssh2 key parsing may fail
      }
    }

    if (!useSSH2) {
      keyType = detectKeyTypeFromContent(cleanKey);

      publicKey = "";
    }

    return {
      privateKey: cleanKey,
      publicKey,
      keyType,
      success: keyType !== "unknown",
      error:
        keyType === "unknown"
          ? getUnsupportedPrivateKeyError(cleanKey)
          : undefined,
    };
  } catch (error) {
    try {
      const cleanKey = normalizePrivateKeyText(privateKeyData);
      const fallbackKeyType = detectKeyTypeFromContent(cleanKey);
      if (fallbackKeyType !== "unknown") {
        return {
          privateKey: cleanKey,
          publicKey: "",
          keyType: fallbackKeyType,
          success: true,
        };
      }
    } catch {
      // expected - fallback key type detection may fail
    }

    const parserError = getErrorMessage(error, "");
    const isPuttyKey = PUTTY_PRIVATE_KEY_RE.test(privateKeyData.trim());

    return {
      privateKey: privateKeyData,
      publicKey: "",
      keyType: "unknown",
      success: false,
      error:
        isPuttyKey && /unsupported|parse|format/i.test(parserError)
          ? getUnsupportedPrivateKeyError(privateKeyData)
          : parserError || getUnsupportedPrivateKeyError(privateKeyData),
    };
  }
}

export function parsePublicKey(publicKeyData: string): PublicKeyInfo {
  try {
    const keyType = detectPublicKeyTypeFromContent(publicKeyData);

    return {
      publicKey: publicKeyData,
      keyType,
      success: keyType !== "unknown",
    };
  } catch (error) {
    return {
      publicKey: publicKeyData,
      keyType: "unknown",
      success: false,
      error: getErrorMessage(error, "Unknown error parsing public key"),
    };
  }
}

export function getFriendlyKeyTypeName(keyType: string): string {
  const keyTypeMap: Record<string, string> = {
    "ssh-rsa": "RSA",
    "ssh-ed25519": "Ed25519",
    "ecdsa-sha2-nistp256": "ECDSA P-256",
    "ecdsa-sha2-nistp384": "ECDSA P-384",
    "ecdsa-sha2-nistp521": "ECDSA P-521",
    "ssh-dss": "DSA",
    "rsa-sha2-256": "RSA-SHA2-256",
    "rsa-sha2-512": "RSA-SHA2-512",
    unknown: "Unknown",
  };

  return keyTypeMap[keyType] || keyType;
}

export function validateKeyPair(
  privateKeyData: string,
  publicKeyData: string,
  passphrase?: string,
): KeyPairValidationResult {
  try {
    const privateKeyInfo = parseSSHKey(privateKeyData, passphrase);
    const publicKeyInfo = parsePublicKey(publicKeyData);

    if (!privateKeyInfo.success) {
      return {
        isValid: false,
        privateKeyType: privateKeyInfo.keyType,
        publicKeyType: publicKeyInfo.keyType,
        error: `Invalid private key: ${privateKeyInfo.error}`,
      };
    }

    if (!publicKeyInfo.success) {
      return {
        isValid: false,
        privateKeyType: privateKeyInfo.keyType,
        publicKeyType: publicKeyInfo.keyType,
        error: `Invalid public key: ${publicKeyInfo.error}`,
      };
    }

    if (privateKeyInfo.keyType !== publicKeyInfo.keyType) {
      return {
        isValid: false,
        privateKeyType: privateKeyInfo.keyType,
        publicKeyType: publicKeyInfo.keyType,
        error: `Key type mismatch: private key is ${privateKeyInfo.keyType}, public key is ${publicKeyInfo.keyType}`,
      };
    }

    if (privateKeyInfo.publicKey && privateKeyInfo.publicKey.trim()) {
      const generatedPublicKey = privateKeyInfo.publicKey.trim();
      const providedPublicKey = publicKeyData.trim();

      const generatedKeyParts = generatedPublicKey.split(" ");
      const providedKeyParts = providedPublicKey.split(" ");

      if (generatedKeyParts.length >= 2 && providedKeyParts.length >= 2) {
        const generatedKeyData =
          generatedKeyParts[0] + " " + generatedKeyParts[1];
        const providedKeyData = providedKeyParts[0] + " " + providedKeyParts[1];

        if (generatedKeyData === providedKeyData) {
          return {
            isValid: true,
            privateKeyType: privateKeyInfo.keyType,
            publicKeyType: publicKeyInfo.keyType,
            generatedPublicKey: generatedPublicKey,
          };
        } else {
          return {
            isValid: false,
            privateKeyType: privateKeyInfo.keyType,
            publicKeyType: publicKeyInfo.keyType,
            generatedPublicKey: generatedPublicKey,
            error: "Public key does not match the private key",
          };
        }
      }
    }

    return {
      isValid: true,
      privateKeyType: privateKeyInfo.keyType,
      publicKeyType: publicKeyInfo.keyType,
      error: "Unable to verify key pair match, but key types are compatible",
    };
  } catch (error) {
    return {
      isValid: false,
      privateKeyType: "unknown",
      publicKeyType: "unknown",
      error: getErrorMessage(error, "Unknown error during validation"),
    };
  }
}