feat: ML-KEM-768, ML-KEM-1024 generate and recover

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<CFError>? = 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")