SSH Key Generator — Swift source

// ssh-keygen — Ed25519 / RSA SSH key pair generation in OpenSSH format.
//
// Language: Swift 5.9+ (Foundation + CryptoKit + Security)
// Ported from src/lib/ssh-keygen.ts
// display source — part of CosmoDev's polyglot tool pages
//
// Keys are created on-device and never leave it. Formats produced:
// - Ed25519 public key : `ssh-ed25519 <base64(blob)> <comment>` where blob is
//   the OpenSSH wire format (string "ssh-ed25519" + string 32-byte key).
// - Ed25519 private key: `-----BEGIN OPENSSH PRIVATE KEY-----` PEM containing
//   the unencrypted "openssh-key-v1" structure (cipher "none", kdf "none").
// - RSA public key     : `ssh-rsa <base64(blob)> <comment>` where blob is
//   string "ssh-rsa" + mpint e + mpint n.
// - RSA private key    : PKCS#8 PEM (`-----BEGIN PRIVATE KEY-----`).
// - Fingerprint        : SHA-256 of the raw public key blob, base64-encoded
//   without padding, prefixed "SHA256:" (matches OpenSSH).

import Foundation
import CryptoKit
import Security

// MARK: - Types

enum SSHAlgorithm: String {
    case ed25519
    case rsa2048 = "rsa-2048"
    case rsa4096 = "rsa-4096"
}

struct SSHKeyOptions {
    var algorithm: SSHAlgorithm
    /// Optional key comment, conventionally `user@host`. Omitted when empty.
    var comment: String? = nil
}

struct SSHKeyPair {
    let publicKey: String
    let privateKey: String
    let fingerprint: String
}

enum SSHKeygenError: Error, CustomStringConvertible {
    case unsupportedAlgorithm(String)
    case badEd25519PublicLength(Int)
    case badEd25519PrivateEncoding
    case rsaKeygen(String)
    case rsaExport
    case rsaMultipartParse

    var description: String {
        switch self {
        case .unsupportedAlgorithm(let a): return "Unsupported algorithm: \(a)"
        case .badEd25519PublicLength(let n): return "Unexpected Ed25519 public key length: \(n)."
        case .badEd25519PrivateEncoding: return "Unexpected Ed25519 private key encoding."
        case .rsaKeygen(let why): return "RSA key generation failed: \(why)"
        case .rsaExport: return "Failed to export the RSA key."
        case .rsaMultipartParse: return "Failed to parse the RSA public key numbers."
        }
    }
}

// MARK: - Wire-format helpers

let opensshMagic: [UInt8] = [
    0x6f, 0x70, 0x65, 0x6e, 0x73, 0x73, 0x68, 0x2d, 0x6b, 0x65, 0x79, 0x2d, 0x76, 0x31, 0x00,
] // "openssh-key-v1\0"

func concat(_ parts: [[UInt8]]) -> [UInt8] {
    parts.reduce(into: []) { out, p in out += p }
}

func u32(_ n: Int) -> [UInt8] {
    [UInt8((n >> 24) & 0xff), UInt8((n >> 16) & 0xff), UInt8((n >> 8) & 0xff), UInt8(n & 0xff)]
}

/// SSH "string": 4-byte big-endian length + raw bytes.
func sshString(_ data: [UInt8]) -> [UInt8] {
    concat([u32(data.count), data])
}

func sshString(_ str: String) -> [UInt8] {
    sshString(Array(str.utf8))
}

/// SSH "mpint": big-endian integer, minimal, with a leading zero when the high bit is set.
func mpint(_ bytes: [UInt8]) -> [UInt8] {
    var i = 0
    while i < bytes.count - 1 && bytes[i] == 0 { i += 1 }
    let v = Array(bytes[i...])
    if !v.isEmpty && v[0] & 0x80 != 0 {
        return sshString(concat([[0], v]))
    }
    return sshString(v)
}

func toBase64(_ bytes: [UInt8]) -> String {
    Data(bytes).base64EncodedString()
}

func pemWrap(_ base64: String, label: String, width: Int) -> String {
    var lines: [String] = []
    var idx = base64.startIndex
    while idx < base64.endIndex {
        let end = base64.index(idx, offsetBy: width, limitedBy: base64.endIndex) ?? base64.endIndex
        lines.append(String(base64[idx..<end]))
        idx = end
    }
    return "-----BEGIN \(label)-----\n\(lines.joined(separator: "\n"))\n-----END \(label)-----\n"
}

/// "SHA256:<base64 digest of the public blob, unpadded>" - the OpenSSH fingerprint format.
func sshFingerprint(_ pubBlob: [UInt8]) -> String {
    let digest = Data(SHA256.hash(data: Data(pubBlob))).base64EncodedString()
    return "SHA256:\(digest.replacingOccurrences(of: "=", with: ""))"
}

// MARK: - Ed25519 private key PEM

/**
 * Build the unencrypted "openssh-key-v1" private key structure for Ed25519
 * (cipher "none", kdf "none"), wrapped in an OPENSSH PRIVATE KEY PEM.
 */
func openSshPrivatePem(_ pubBlob: [UInt8], _ privateKeyBytes: [UInt8], _ comment: String) -> String {
    var rng = SystemRandomNumberGenerator()
    let checkInt = UInt32.random(in: .min ... .max, using: &rng)
    let check = u32(Int(checkInt))
    var inner = concat([check, check, pubBlob, sshString(privateKeyBytes), sshString(comment)])
    // Pad to a multiple of the block size (8 for cipher "none") with 1,2,3,...
    let padLen = (8 - (inner.count % 8)) % 8
    inner += (1...padLen).map(UInt8.init)
    let outer = concat([
        opensshMagic,
        sshString("none"),        // ciphername
        sshString("none"),        // kdfname
        sshString([]),            // kdfoptions
        u32(1),                   // number of keys
        sshString(pubBlob),
        sshString(inner),
    ])
    return pemWrap(toBase64(outer), label: "OPENSSH PRIVATE KEY", width: 70)
}

// MARK: - RSA helpers (Security framework - CryptoKit has no RSA)

/// Read one TLV header at `pos`: returns (tag, content, nextPosition).
private func derReadTLV(_ bytes: [UInt8], _ pos: Int) -> (tag: UInt8, content: ArraySlice<UInt8>, next: Int)? {
    guard pos + 2 <= bytes.count else { return nil }
    let tag = bytes[pos]
    let first = bytes[pos + 1]
    var len = 0
    var headerEnd = pos + 2
    if first < 0x80 {
        len = Int(first)
    } else {
        let numBytes = Int(first & 0x7f)
        guard numBytes > 0, numBytes <= 4, pos + 2 + numBytes <= bytes.count else { return nil }
        for i in 0..<numBytes { len = (len << 8) | Int(bytes[pos + 2 + i]) }
        headerEnd = pos + 2 + numBytes
    }
    guard headerEnd + len <= bytes.count else { return nil }
    return (tag, bytes[headerEnd..<(headerEnd + len)], headerEnd + len)
}

/// Extract (modulus n, exponent e) from a PKCS#1 RSAPublicKey DER blob.
func rsaPublicNumbers(fromPKCS1 pkcs1: [UInt8]) throws -> (n: [UInt8], e: [UInt8]) {
    guard let seq = derReadTLV(pkcs1, 0), seq.tag == 0x30 else { throw SSHKeygenError.rsaMultipartParse }
    let body = Array(seq.content)
    guard let modTLV = derReadTLV(body, 0), modTLV.tag == 0x02,
          let expTLV = derReadTLV(body, modTLV.next), expTLV.tag == 0x02 else {
        throw SSHKeygenError.rsaMultipartParse
    }
    return (Array(modTLV.content), Array(expTLV.content))
}

// MARK: - Generation

/// Generate an SSH key pair in OpenSSH format (CryptoKit + Security).
func generateSSHKeyPair(_ options: SSHKeyOptions) throws -> SSHKeyPair {
    let comment = options.comment?.trimmingCharacters(in: .whitespaces) ?? ""
    let suffix = comment.isEmpty ? "" : " \(comment)"

    switch options.algorithm {
    case .ed25519:
        let key = Curve25519.Signing.PrivateKey()
        let pub = [UInt8](key.publicKey.rawRepresentation)
        guard pub.count == 32 else { throw SSHKeygenError.badEd25519PublicLength(pub.count) }
        // rawRepresentation IS the 32-byte seed for Curve25519 private keys.
        let seed = [UInt8](key.rawRepresentation)
        guard seed.count == 32 else { throw SSHKeygenError.badEd25519PrivateEncoding }
        let pubBlob = concat([sshString("ssh-ed25519"), sshString(pub)])
        return SSHKeyPair(
            publicKey: "ssh-ed25519 \(toBase64(pubBlob))\(suffix)",
            privateKey: openSshPrivatePem(pubBlob, concat([seed, pub]), comment),
            fingerprint: sshFingerprint(pubBlob)
        )

    case .rsa2048, .rsa4096:
        let bits = options.algorithm == .rsa4096 ? 4096 : 2048
        var error: Unmanaged<CFError>?
        let params: [String: Any] = [
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecAttrKeySizeInBits as String: bits,
        ]
        guard let secKey = SecKeyCreateRandomKey(params as CFDictionary, &error) else {
            throw SSHKeygenError.rsaKeygen(String(describing: error?.takeRetainedValue()))
        }
        guard let publicKey = SecKeyCopyPublicKey(secKey),
              let pkcs1 = SecKeyCopyExternalRepresentation(publicKey, &error) as Data? else {
            throw SSHKeygenError.rsaExport
        }
        // Public: SecKey exports PKCS#1 (RSAPublicKey) - parse out e and n.
        let numbers = try rsaPublicNumbers(fromPKCS1: [UInt8](pkcs1))
        let pubBlob = concat([
            sshString("ssh-rsa"),
            mpint(numbers.e),
            mpint(numbers.n),
        ])
        // Private: SecKey exports RSA private keys as PKCS#8 directly.
        guard let pkcs8 = SecKeyCopyExternalRepresentation(secKey, &error) as Data? else {
            throw SSHKeygenError.rsaExport
        }
        return SSHKeyPair(
            publicKey: "ssh-rsa \(toBase64(pubBlob))\(suffix)",
            privateKey: pemWrap(toBase64([UInt8](pkcs8)), label: "PRIVATE KEY", width: 64),
            fingerprint: sshFingerprint(pubBlob)
        )
    }
}