From a652fff5018def2e63f25349429a2ac2ceded3fe Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Sat, 27 Sep 2025 19:28:40 +0800 Subject: [PATCH] feat: ML-KEM-768, ML-KEM-1024 ecdh --- swift-secure-enclave-tool-v2.swift | 111 +++++++++++++++++++++++------ 1 file changed, 89 insertions(+), 22 deletions(-) diff --git a/swift-secure-enclave-tool-v2.swift b/swift-secure-enclave-tool-v2.swift index 686849d..a276c43 100644 --- a/swift-secure-enclave-tool-v2.swift +++ b/swift-secure-enclave-tool-v2.swift @@ -9,6 +9,12 @@ import CryptoKit import Foundation import LocalAuthentication +extension Data { + var hexEncodedString: String { + return map { String(format: "%02x", $0) }.joined() + } +} + struct GenerateSecureEnclaveKeyPairRequest { var controlFlag: String } @@ -131,12 +137,12 @@ func parseExternalSignRequest() -> ComputeP256EcSignRequest? { ) } -struct ComputeP256EcdhRequest { +struct ComputeEcdhRequest { var dataRepresentationBase64: String var ephemeraPublicKeyBase64: String } -func parseComputeP256EcdhRequest() -> ComputeP256EcdhRequest? { +func parseComputeEcdhRequest() -> ComputeEcdhRequest? { var dataRepresentationBase64Opt: String? var ephemeraPublicKeyBase64Opt: String? let len = CommandLine.arguments.count; @@ -163,7 +169,7 @@ func parseComputeP256EcdhRequest() -> ComputeP256EcdhRequest? { exitError("parameter --ephemera-public-key-base64 required.") return nil } - return ComputeP256EcdhRequest( + return ComputeEcdhRequest( dataRepresentationBase64: dataRepresentationBase64, ephemeraPublicKeyBase64: ephemeraPublicKeyBase64 ) @@ -259,7 +265,7 @@ struct ComputeSecureEnclaveP256EcsignResponse: Codable { var signature_base64: String } -struct ComputeSecureEnclaveP256EcdhResponse: Codable { +struct ComputeSecureEnclaveEcdhResponse: Codable { var success: Bool var shared_secret_hex: String } @@ -580,7 +586,7 @@ func computeSecureEnclaveP256Ecsign(request: ComputeP256EcSignRequest) -> Comput } } -func computeSecureEnclaveP256Ecdh(request: ComputeP256EcdhRequest) -> ComputeSecureEnclaveP256EcdhResponse? { +func computeSecureEnclaveP256Ecdh(request: ComputeEcdhRequest) -> ComputeSecureEnclaveEcdhResponse? { guard let privateKeyDataRepresentation = Data( base64Encoded: request.dataRepresentationBase64 ) else { @@ -614,7 +620,7 @@ func computeSecureEnclaveP256Ecdh(request: ComputeP256EcdhRequest) -> ComputeSec shared_secret_hex = sharedSecret.description } - return ComputeSecureEnclaveP256EcdhResponse( + return ComputeSecureEnclaveEcdhResponse( success: true, shared_secret_hex: shared_secret_hex ) @@ -624,6 +630,53 @@ func computeSecureEnclaveP256Ecdh(request: ComputeP256EcdhRequest) -> ComputeSec } } +func computeSecureEnclaveMlKemEcdh(request: ComputeEcdhRequest, keyLen: Int) -> ComputeSecureEnclaveEcdhResponse? { + guard let privateKeyDataRepresentation = Data( + base64Encoded: request.dataRepresentationBase64 + ) else { + exitError("private key base64 decode failed") + return nil + } + guard let ephemeralPublicKeyRepresentation = Data( + base64Encoded: request.ephemeraPublicKeyBase64 + ) else { + exitError("ephemeral public key base64 decode failed") + return nil + } + do { + let context = LAContext() + if keyLen == 768 { + let p = try SecureEnclave.MLKEM768.PrivateKey( + dataRepresentation: privateKeyDataRepresentation, + authenticationContext: context + ) + let sharedSecret = try p.decapsulate(ephemeralPublicKeyRepresentation) + let sharedSecretData = sharedSecret.withUnsafeBytes { Data($0) } + return ComputeSecureEnclaveEcdhResponse( + success: true, + shared_secret_hex: sharedSecretData.hexEncodedString + ) + } else if (keyLen == 1024) { + let p = try SecureEnclave.MLKEM1024.PrivateKey( + dataRepresentation: privateKeyDataRepresentation, + authenticationContext: context + ) + let sharedSecret = try p.decapsulate(ephemeralPublicKeyRepresentation) + let sharedSecretData = sharedSecret.withUnsafeBytes { Data($0) } + return ComputeSecureEnclaveEcdhResponse( + success: true, + shared_secret_hex: sharedSecretData.hexEncodedString + ) + } else { + exitError("Invalid algorithm: ML-KEM-\(keyLen)") + return nil + } + } catch { + exitError("\(error)") + return nil + } +} + func externalSpec() -> ExternalSpecResponse { return ExternalSpecResponse( success: true, @@ -694,11 +747,23 @@ if (command == "compute_p256_ecsign") { } if (command == "compute_p256_ecdh") { - let request = parseComputeP256EcdhRequest()!; + let request = parseComputeEcdhRequest()!; let response = computeSecureEnclaveP256Ecdh(request: request) exitOkWithJson(response) } +if (command == "compute_mlkem768_ecdh") { + let request = parseComputeEcdhRequest()!; + let response = computeSecureEnclaveMlKemEcdh(request: request, keyLen: 768) + exitOkWithJson(response) +} + +if (command == "compute_mlkem1024_ecdh") { + let request = parseComputeEcdhRequest()!; + let response = computeSecureEnclaveMlKemEcdh(request: request, keyLen: 1024) + exitOkWithJson(response) +} + if (command == "external_spec") { exitOkWithJson(externalSpec()) } @@ -721,22 +786,24 @@ if (command == "version") { if (command == "help" || command == "-h" || command == "--help") { print("swift-secure-enclave-tool-v2 [parameters]") - print("help - print help") - print("version - print version") - print("is_support_secure_enclave - is Secure Enclave supported") - print("generate_p256_ecsign_keypair --control-flag <> - generate Secure Enclave P256 EC sign key pair") - print("generate_p256_ecdh_keypair --control-flag <> - generate Secure Enclave P256 EC DH key pair") - print("generate_mlkem768_ecdh_keypair --control-flag <> - generate Secure Enclave ML-KEM-768 key pair") - print("generate_mlkem1024_ecdh_keypair --control-flag <> - generate Secure Enclave ML-KEM-1024 key pair") - print("recover_p256_ecsign_public_key --private-key <> - recover Secure Enclave P256 EC sign key pair") - print("recover_p256_ecdh_public_key --private-key <> - recover Secure Enclave P256 EC DH key pair") - print("recover_mlkem768_public_key --private-key <> - recover Secure Enclave ML-KEM-768 key pair") - print("recover_mlkem1024_public_key --private-key <> - recover Secure Enclave ML-KEM-1024 key pair") + print("help - print help") + print("version - print version") + print("is_support_secure_enclave - is Secure Enclave supported") + print("generate_p256_ecsign_keypair --control-flag <> - generate Secure Enclave P256 EC sign key pair") + print("generate_p256_ecdh_keypair --control-flag <> - generate Secure Enclave P256 EC DH key pair") + print("generate_mlkem768_ecdh_keypair --control-flag <> - generate Secure Enclave ML-KEM-768 key pair") + print("generate_mlkem1024_ecdh_keypair --control-flag <> - generate Secure Enclave ML-KEM-1024 key pair") + print("recover_p256_ecsign_public_key --private-key <> - recover Secure Enclave P256 EC sign key pair") + print("recover_p256_ecdh_public_key --private-key <> - recover Secure Enclave P256 EC DH key pair") + print("recover_mlkem768_public_key --private-key <> - recover Secure Enclave ML-KEM-768 key pair") + print("recover_mlkem1024_public_key --private-key <> - recover Secure Enclave ML-KEM-1024 key pair") print("compute_p256_ecsign --private-key <> --message-base64 <> [--message-type <>] - compure Secure Enclave P256 EC sign") - print("compute_p256_ecdh --private-key <> --ephemera-public-key <> - compure Secure Enclave P256 EC DH") - print("external_spec - external specification") - print("external_public_key --parameter <> - external public key") - print("external_sign --parameter <> --alg ES256 --message-base64 <> - external sign") + print("compute_p256_ecdh --private-key <> --ephemera-public-key <> - compure Secure Enclave P256 EC DH") + print("compute_mlkem768_ecdh --private-key <> --ephemera-public-key <> - compure Secure Enclave ML-KEM-768") + print("compute_mlkem1024_ecdh --private-key <> --ephemera-public-key <> - compure Secure Enclave ML-KEM-1024") + print("external_spec - external specification") + print("external_public_key --parameter <> - external public key") + print("external_sign --parameter <> --alg ES256 --message-base64 <> - external sign") print() print("options:") print("> --control-flag - none, userPresence, devicePasscode, biometryAny, biometryCurrentSet")