WireGuard Key Generator — Swift source

// wireguard-keygen — WireGuard key generation (Curve25519, RFC 7748).
//
// Language: Swift 5.9+ (Foundation + CryptoKit)
// Ported from src/lib/wireguard-keygen.ts
// display source — part of CosmoDev's polyglot tool pages
//
// WireGuard uses Curve25519 (RFC 7748 X25519) for its key exchange:
//   - a private key is 32 random bytes, clamped per the Curve25519 rules
//     (`key[0] &= 248; key[31] &= 127; key[31] |= 64`)
//   - the public key is that scalar multiplied by the curve's base point 9
//   - an optional pre-shared key is 32 random bytes, used as-is (no clamping)
//
// The TS reference implements the X25519 Montgomery ladder with plain BigInt
// arithmetic because JS ships no curve primitive. CryptoKit ships X25519
// natively (Curve25519.KeyAgreement), so this port computes the same scalar
// product through it. The clamp function is kept verbatim: CryptoKit clamps
// internally too, but clamping before serialization makes the generated
// private key bytes identical to the TS reference's.
//
// Every key is standard padded Base64 — 44 characters for 32 bytes — exactly
// what WireGuard config files expect. Randomness comes from
// SystemRandomNumberGenerator, the system CSPRNG.

import Foundation
import CryptoKit

// MARK: - Types

/// A WireGuard key pair, both keys in standard padded Base64.
struct WireGuardKeys: Equatable {
    let privateKey: String
    let publicKey: String
}

enum WireGuardError: Error, LocalizedError {
    case invalidBase64Length
    case invalidBase64Character(Character)
    case wrongKeyLength(expected: Int, got: Int)

    var errorDescription: String? {
        switch self {
        case .invalidBase64Length:
            return "Invalid Base64: length must be a non-zero multiple of 4"
        case .invalidBase64Character(let ch):
            return "Invalid Base64 character: \"\(ch)\""
        case .wrongKeyLength(let expected, let got):
            let what = expected == KEY_LENGTH ? "Private key" : "Key"
            return "Invalid private key: expected \(expected) bytes, got \(got) (\(what))"
        }
    }
}

/// Length of every WireGuard key, in bytes.
let KEY_LENGTH = 32

/// The curve's base point u = 9, little-endian — what every private key is
/// multiplied by to derive its public key.
let BASE_POINT: [UInt8] = {
    var u = [UInt8](repeating: 0, count: KEY_LENGTH)
    u[0] = 9
    return u
}()

// MARK: - Base64 codec (standard alphabet, always padded — the WireGuard format)

let B64_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

let B64_VALUES: [Character: Int] = {
    var map: [Character: Int] = [:]
    for (i, c) in B64_ALPHABET.enumerated() { map[c] = i }
    return map
}()

/// Encode bytes as standard Base64 with `=` padding (32 bytes → 44 chars).
/// (Foundation's Data base64 encoding is exactly this format.)
func bytesToBase64(_ bytes: [UInt8]) -> String {
    Data(bytes).base64EncodedString()
}

/// Decode standard Base64 (with padding). Throws on invalid input.
func base64ToBytes(_ b64: String) throws -> [UInt8] {
    let s = b64.trimmingCharacters(in: .whitespacesAndNewlines)
    guard !s.isEmpty, s.count % 4 == 0 else { throw WireGuardError.invalidBase64Length }
    let pad = s.hasSuffix("==") ? 2 : (s.hasSuffix("=") ? 1 : 0)
    var out = [UInt8]()
    out.reserveCapacity(s.count / 4 * 3 - pad)
    var buffer = 0
    var bits = 0
    for ch in s.dropLast(pad) {
        guard let v = B64_VALUES[ch] else { throw WireGuardError.invalidBase64Character(ch) }
        buffer = (buffer << 6) | v
        bits += 6
        if bits >= 8 {
            bits -= 8
            out.append(UInt8((buffer >> bits) & 0xff))
        }
    }
    return out
}

// MARK: - Curve25519

/// Clamp 32 bytes into a valid Curve25519 scalar (RFC 7748 §5). Returns a copy.
func clampPrivateKey(_ key: [UInt8]) throws -> [UInt8] {
    guard key.count == KEY_LENGTH else {
        throw WireGuardError.wrongKeyLength(expected: KEY_LENGTH, got: key.count)
    }
    var out = key
    out[0] &= 248   // clear the low 3 bits → multiple of the cofactor
    out[31] &= 127  // clear the high bit
    out[31] |= 64   // force bit 254 → the ladder always sees a 255-bit scalar
    return out
}

/// X25519 scalar multiplication `scalar · u` — the same operation as the TS
/// reference's BigInt Montgomery ladder, computed by CryptoKit's native
/// implementation. The scalar is clamped internally (an unclamped input
/// yields the same result as its clamped form, exactly like every X25519
/// implementation).
func curve25519(scalar: [UInt8], u: [UInt8]) throws -> [UInt8] {
    guard scalar.count == KEY_LENGTH else {
        throw WireGuardError.wrongKeyLength(expected: KEY_LENGTH, got: scalar.count)
    }
    guard u.count == KEY_LENGTH else {
        throw WireGuardError.wrongKeyLength(expected: KEY_LENGTH, got: u.count)
    }
    let secret = try Curve25519.KeyAgreement.PrivateKey(rawRepresentation: Data(scalar))
    let base = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: Data(u))
    let shared = try secret.sharedSecretFromKeyAgreement(with: base)
    return Array(shared.withUnsafeBytes { Data($0) })
}

// MARK: - Key generation

/// 32 CSPRNG bytes — SystemRandomNumberGenerator is the system CSPRNG
/// (SecRandom on Apple platforms), the counterpart of crypto.getRandomValues.
func randomBytes(_ length: Int) -> [UInt8] {
    var csprng = SystemRandomNumberGenerator()
    return (0..<length).map { _ in UInt8.random(in: .min ... .max, using: &csprng) }
}

/// A fresh private key: 32 CSPRNG bytes, clamped, Base64.
func generatePrivateKey() -> String {
    (try? bytesToBase64(clampPrivateKey(randomBytes(KEY_LENGTH)))) ?? ""
}

/// A fresh pre-shared key: 32 CSPRNG bytes, Base64 — used as-is, never clamped.
func generatePresharedKey() -> String {
    bytesToBase64(randomBytes(KEY_LENGTH))
}

/// Derive the WireGuard public key that pairs with a Base64 private key
/// (Curve25519 scalar multiplication of the base point).
func privateKeyToPublic(_ privateKeyBase64: String) throws -> String {
    let priv = try base64ToBytes(privateKeyBase64)
    guard priv.count == KEY_LENGTH else {
        throw WireGuardError.wrongKeyLength(expected: KEY_LENGTH, got: priv.count)
    }
    return bytesToBase64(try curve25519(scalar: priv, u: BASE_POINT))
}

/// A fresh WireGuard key pair (private + matching public key, both Base64).
func generateWireGuardKeys() throws -> WireGuardKeys {
    let privateKey = generatePrivateKey()
    let publicKey = try privateKeyToPublic(privateKey)
    return WireGuardKeys(privateKey: privateKey, publicKey: publicKey)
}

// MARK: - Config template

/// Render a `wg-quick` config template around a key pair. The peer's public
/// key, endpoint, and your tunnel address depend on the other side, so they
/// stay as placeholders. A `PresharedKey` line is included only when `psk` is
/// given (it must be present on BOTH sides of the tunnel).
func formatConfig(_ keys: WireGuardKeys, psk: String? = nil) -> String {
    var lines = [
        "[Interface]",
        "# Your side — keep PrivateKey secret, share PublicKey with the peer",
        "PrivateKey = \(keys.privateKey)",
        "PublicKey = \(keys.publicKey)",
        "# Tunnel address assigned by your server (plus optional tunnel DNS)",
        "Address = 10.0.0.2/32",
        "# DNS = 1.1.1.1",
        "",
        "[Peer]",
        "# The other side's public key",
        "PublicKey = <PEER_PUBLIC_KEY>",
    ]
    if let psk = psk, !psk.isEmpty {
        lines.append("# Optional pre-shared key — the same value must be set on BOTH sides")
        lines.append("PresharedKey = \(psk)")
    }
    lines += [
        "# Route everything through the tunnel (or scope it, e.g. 10.0.0.0/24)",
        "AllowedIPs = 0.0.0.0/0, ::/0",
        "# The peer's public address and port",
        "Endpoint = vpn.example.com:51820",
        "PersistentKeepalive = 25",
    ]
    return lines.joined(separator: "\n")
}