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

import ssh2Pkg from "ssh2";
const ssh2Utils = ssh2Pkg.utils;

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;

export 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.";
}

export function preparePrivateKeyForSSH2(
  privateKeyData: string,
  passphrase?: string,
): Buffer {
  const cleanKey = normalizePrivateKeyText(privateKeyData);
  const keyInfo = parseSSHKey(cleanKey, passphrase);

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

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

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 = error instanceof Error ? error.message : "";
    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:
        error instanceof Error
          ? error.message
          : "Unknown error parsing public key",
    };
  }
}

export function detectKeyType(privateKeyData: string): string {
  try {
    const parsedKey = ssh2Utils.parseKey(privateKeyData);
    if (parsedKey instanceof Error) {
      return "unknown";
    }
    return parsedKey.type || "unknown";
  } catch {
    return "unknown";
  }
}

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:
        error instanceof Error
          ? error.message
          : "Unknown error during validation",
    };
  }
}