From f310a5f26438c50e9b5508e6135c6fec7ae2a5cc Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Sat, 27 Sep 2025 19:12:58 +0800 Subject: [PATCH] feat: ML-KEM-768, ML-KEM-1024 generate and recover --- swift-secure-enclave-tool-v2.swift | 146 +++++++++++++++++++++++++---- 1 file changed, 130 insertions(+), 16 deletions(-) diff --git a/swift-secure-enclave-tool-v2.swift b/swift-secure-enclave-tool-v2.swift index 8887e72..686849d 100644 --- a/swift-secure-enclave-tool-v2.swift +++ b/swift-secure-enclave-tool-v2.swift @@ -9,11 +9,11 @@ import CryptoKit import Foundation import LocalAuthentication -struct GenerateSecureEnclaveP256KeyPairRequest { +struct GenerateSecureEnclaveKeyPairRequest { var controlFlag: String } -func parseGenerateSecureEnclaveP256KeyPairRequest() -> GenerateSecureEnclaveP256KeyPairRequest? { +func parseGenerateSecureEnclaveKeyPairRequest() -> GenerateSecureEnclaveKeyPairRequest? { var controlFlagOpt: String? let len = CommandLine.arguments.count; if CommandLine.arguments.count > 2 { @@ -32,7 +32,7 @@ func parseGenerateSecureEnclaveP256KeyPairRequest() -> GenerateSecureEnclaveP256 exitError("parameter --control-flag required.") return nil } - return GenerateSecureEnclaveP256KeyPairRequest( + return GenerateSecureEnclaveKeyPairRequest( controlFlag: controlFlag ) } @@ -169,11 +169,11 @@ func parseComputeP256EcdhRequest() -> ComputeP256EcdhRequest? { ) } -struct RecoverSecureEnclaveP256PublicKeyRequest { +struct RecoverSecureEnclavePublicKeyRequest { var dataRepresentationBase64: String } -func parseRecoverSecureEnclaveP256PublicKeyRequest() -> RecoverSecureEnclaveP256PublicKeyRequest? { +func parseRecoverSecureEnclavePublicKeyRequest() -> RecoverSecureEnclavePublicKeyRequest? { var dataRepresentationBase64Opt: String? let len = CommandLine.arguments.count; if CommandLine.arguments.count > 2 { @@ -192,12 +192,12 @@ func parseRecoverSecureEnclaveP256PublicKeyRequest() -> RecoverSecureEnclaveP256 exitError("parameter --data-representation-base64 or --private-key required.") return nil } - return RecoverSecureEnclaveP256PublicKeyRequest( + return RecoverSecureEnclavePublicKeyRequest( dataRepresentationBase64: dataRepresentationBase64 ) } -func parseExternalPublicKeyRequest() -> RecoverSecureEnclaveP256PublicKeyRequest? { +func parseExternalPublicKeyRequest() -> RecoverSecureEnclavePublicKeyRequest? { var dataRepresentationBase64Opt: String? let len = CommandLine.arguments.count; if CommandLine.arguments.count > 2 { @@ -216,7 +216,7 @@ func parseExternalPublicKeyRequest() -> RecoverSecureEnclaveP256PublicKeyRequest exitError("parameter --parameter required.") return nil } - return RecoverSecureEnclaveP256PublicKeyRequest( + return RecoverSecureEnclavePublicKeyRequest( dataRepresentationBase64: dataRepresentationBase64 ) } @@ -243,6 +243,12 @@ struct GenerateSecureEnclaveP256KeyPairResponse: Codable { var data_representation_base64: String } +struct GenerateSecureEnclaveMlKemKeyPairResponse: Codable { + var success: Bool + var public_key_base64: String + var data_representation_base64: String +} + struct ExternalPublicKeyResponse: Codable { var success: Bool var public_key_base64: String @@ -347,10 +353,9 @@ func isSupportSecureEnclave() -> SupportSecureEnclaveResponse { return SupportSecureEnclaveResponse(success: true, supported: SecureEnclave.isAvailable) } -func generateSecureEnclaveP256KeyPair(sign: Bool, request: GenerateSecureEnclaveP256KeyPairRequest) -> GenerateSecureEnclaveP256KeyPairResponse? { +func getSecAccessControlCreateWithFlags(controlFlag: String) -> SecAccessControl? { var error: Unmanaged? = nil let accessControlCreateFlags: SecAccessControlCreateFlags - let controlFlag = request.controlFlag if (controlFlag == "none") { accessControlCreateFlags = [.privateKeyUsage] } else if (controlFlag == "userPresence") { @@ -374,6 +379,13 @@ func generateSecureEnclaveP256KeyPair(sign: Bool, request: GenerateSecureEnclave exitError(error.debugDescription) return nil } + return accessCtrl +} + +func generateSecureEnclaveP256KeyPair(sign: Bool, request: GenerateSecureEnclaveKeyPairRequest) -> GenerateSecureEnclaveP256KeyPairResponse? { + guard let accessCtrl = getSecAccessControlCreateWithFlags(controlFlag: request.controlFlag) else { + return nil + } do { if (sign) { let privateKeyReference = try SecureEnclave.P256.Signing.PrivateKey.init( @@ -392,6 +404,31 @@ func generateSecureEnclaveP256KeyPair(sign: Bool, request: GenerateSecureEnclave } } +func generateSecureEnclaveMlKemKeyPair(keyLen: Int, request: GenerateSecureEnclaveKeyPairRequest) -> GenerateSecureEnclaveMlKemKeyPairResponse? { + guard let accessCtrl = getSecAccessControlCreateWithFlags(controlFlag: request.controlFlag) else { + return nil + } + do { + if (keyLen == 768) { + let privateKeyReference = try CryptoKit.SecureEnclave.MLKEM768.PrivateKey.init( + accessControl: accessCtrl + ); + return mlKem768PrivateKeyToResponse(privateKeyReference); + } else if (keyLen == 1024) { + let privateKeyReference = try CryptoKit.SecureEnclave.MLKEM1024.PrivateKey.init( + accessControl: accessCtrl + ); + return mlKem1024PrivateKeyToResponse(privateKeyReference); + } else { + exitError("Invalid algorithm: ML-KEM-\(keyLen)") + return nil + } + } catch { + exitError("\(error)") + return nil + } +} + func signingPrivateKeyToResponse(_ privateKeyReference: SecureEnclave.P256.Signing.PrivateKey) -> GenerateSecureEnclaveP256KeyPairResponse { let publicKeyPointBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString() let publicKeyDerBase64 = privateKeyReference.publicKey.derRepresentation.base64EncodedString() @@ -406,7 +443,7 @@ func signingPrivateKeyToResponse(_ privateKeyReference: SecureEnclave.P256.Signi func keyAgreementPrivateKeyToResponse(_ privateKeyReference: SecureEnclave.P256.KeyAgreement.PrivateKey) -> GenerateSecureEnclaveP256KeyPairResponse { let publicKeyPointBase64 = privateKeyReference.publicKey.x963Representation.base64EncodedString() - let publicKeyDerBase64 = privateKeyReference.publicKey.derRepresentation.base64EncodedString() + let publicKeyDerBase64 = privateKeyReference.publicKey.rawRepresentation.base64EncodedString() let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString() return GenerateSecureEnclaveP256KeyPairResponse( success: true, @@ -416,7 +453,27 @@ func keyAgreementPrivateKeyToResponse(_ privateKeyReference: SecureEnclave.P256. ) } -func recoverSecureEnclaveP256PublicKey(request: RecoverSecureEnclaveP256PublicKeyRequest, sign: Bool) -> GenerateSecureEnclaveP256KeyPairResponse? { +func mlKem768PrivateKeyToResponse(_ privateKeyReference: SecureEnclave.MLKEM768.PrivateKey) -> GenerateSecureEnclaveMlKemKeyPairResponse { + let publicKeyDerBase64 = privateKeyReference.publicKey.rawRepresentation.base64EncodedString() + let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString() + return GenerateSecureEnclaveMlKemKeyPairResponse( + success: true, + public_key_base64: publicKeyDerBase64, + data_representation_base64: dataRepresentationBase64 + ) +} + +func mlKem1024PrivateKeyToResponse(_ privateKeyReference: SecureEnclave.MLKEM1024.PrivateKey) -> GenerateSecureEnclaveMlKemKeyPairResponse { + let publicKeyDerBase64 = privateKeyReference.publicKey.rawRepresentation.base64EncodedString() + let dataRepresentationBase64 = privateKeyReference.dataRepresentation.base64EncodedString() + return GenerateSecureEnclaveMlKemKeyPairResponse( + success: true, + public_key_base64: publicKeyDerBase64, + data_representation_base64: dataRepresentationBase64 + ) +} + +func recoverSecureEnclaveP256PublicKey(request: RecoverSecureEnclavePublicKeyRequest, sign: Bool) -> GenerateSecureEnclaveP256KeyPairResponse? { guard let privateKeyDataRepresentation = Data( base64Encoded: request.dataRepresentationBase64 ) else { @@ -444,6 +501,37 @@ func recoverSecureEnclaveP256PublicKey(request: RecoverSecureEnclaveP256PublicKe } } +func recoverSecureEnclaveMlKemPublicKey(request: RecoverSecureEnclavePublicKeyRequest, keyLen: Int) -> GenerateSecureEnclaveMlKemKeyPairResponse? { + guard let privateKeyDataRepresentation = Data( + base64Encoded: request.dataRepresentationBase64 + ) else { + exitError("private key base64 decode failed") + return nil + } + do { + let context = LAContext() + if (keyLen == 768) { + let privateKeyReference = try SecureEnclave.MLKEM768.PrivateKey( + dataRepresentation: privateKeyDataRepresentation, + authenticationContext: context + ) + return mlKem768PrivateKeyToResponse(privateKeyReference); + } else if (keyLen == 1024) { + let privateKeyReference = try SecureEnclave.MLKEM1024.PrivateKey( + dataRepresentation: privateKeyDataRepresentation, + authenticationContext: context + ) + return mlKem1024PrivateKeyToResponse(privateKeyReference); + } else { + exitError("Invalid algorithm: ML-KEM-\(keyLen)") + return nil + } + } catch { + exitError("\(error)") + return nil + } +} + func computeSecureEnclaveP256Ecsign(request: ComputeP256EcSignRequest) -> ComputeSecureEnclaveP256EcsignResponse? { guard let privateKeyDataRepresentation = Data( base64Encoded: request.dataRepresentationBase64 @@ -556,27 +644,49 @@ if (command == "is_support_secure_enclave") { } if (command == "generate_p256_ecsign_keypair") { - let request = parseGenerateSecureEnclaveP256KeyPairRequest()!; + let request = parseGenerateSecureEnclaveKeyPairRequest()!; exitOkWithJson(generateSecureEnclaveP256KeyPair(sign: true, request: request)) } if (command == "generate_p256_ecdh_keypair") { - let request = parseGenerateSecureEnclaveP256KeyPairRequest()!; + let request = parseGenerateSecureEnclaveKeyPairRequest()!; exitOkWithJson(generateSecureEnclaveP256KeyPair(sign: false, request: request)) } +if (command == "generate_mlkem768_ecdh_keypair") { + let request = parseGenerateSecureEnclaveKeyPairRequest()!; + exitOkWithJson(generateSecureEnclaveMlKemKeyPair(keyLen: 768, request: request)) +} + +if (command == "generate_mlkem1024_ecdh_keypair") { + let request = parseGenerateSecureEnclaveKeyPairRequest()!; + exitOkWithJson(generateSecureEnclaveMlKemKeyPair(keyLen: 1024, request: request)) +} + if (command == "recover_p256_ecsign_public_key") { - let request = parseRecoverSecureEnclaveP256PublicKeyRequest()! + let request = parseRecoverSecureEnclavePublicKeyRequest()! let response = recoverSecureEnclaveP256PublicKey(request: request, sign: true) exitOkWithJson(response) } if (command == "recover_p256_ecdh_public_key") { - let request = parseRecoverSecureEnclaveP256PublicKeyRequest()! + let request = parseRecoverSecureEnclavePublicKeyRequest()! let response = recoverSecureEnclaveP256PublicKey(request: request, sign: false) exitOkWithJson(response) } +if (command == "recover_mlkem768_public_key") { + let request = parseRecoverSecureEnclavePublicKeyRequest()! + let response = recoverSecureEnclaveMlKemPublicKey(request: request, keyLen: 768) + exitOkWithJson(response) +} + +if (command == "recover_mlkem1024_public_key") { + let request = parseRecoverSecureEnclavePublicKeyRequest()! + let response = recoverSecureEnclaveMlKemPublicKey(request: request, keyLen: 1024) + exitOkWithJson(response) +} + if (command == "compute_p256_ecsign") { let request = parseComputeP256EcSignRequest()!; let response = computeSecureEnclaveP256Ecsign(request: request) @@ -616,8 +726,12 @@ if (command == "help" || command == "-h" || command == "--help") { 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")