medmap-ios/medmap/Utilities/KeyPairManager.swift

//
//  KeyPairManager.swift
//  FluxRead
//
//  Created by RSA on 11/03/2026.
//

import Foundation
import Security
import CryptoKit

// MARK: - KeyPairManager
/// Manages generation and secure storage of RSA / EC key pairs using the iOS Secure Enclave & Keychain.
final class KeyPairManager {

    // MARK: - Configuration

    enum KeyAlgorithm {
        case rsa(bits: Int)        // e.g. .rsa(bits: 2048)
        case ecSecureEnclave       // P-256 backed by Secure Enclave (recommended)
        case ecSoftware            // P-256 in software keychain (simulator / older devices)
    }

    struct KeyPairConfig {
        let tag: String            // Unique reverse-DNS tag  e.g. "com.myapp.userkey"
        let algorithm: KeyAlgorithm
        let requiresBiometrics: Bool
    }

    // MARK: - Errors

    enum KeyPairError: LocalizedError {
        case generationFailed(OSStatus)
        case publicKeyExportFailed(OSStatus)
        case keyNotFound
        case deletionFailed(OSStatus)
        case unsupportedAlgorithm
        case secureEnclaveUnavailable

        var errorDescription: String? {
            switch self {
            case .generationFailed(let s):      return "Key generation failed: \(s)"
            case .publicKeyExportFailed(let s): return "Public key export failed: \(s)"
            case .keyNotFound:                  return "Key not found in Keychain"
            case .deletionFailed(let s):        return "Key deletion failed: \(s)"
            case .unsupportedAlgorithm:         return "Algorithm not supported on this device"
            case .secureEnclaveUnavailable:     return "Secure Enclave is not available"
            }
        }
    }

    // MARK: - Public API

    /// Generate a key pair and store the **private key** securely in the Keychain / Secure Enclave.
    /// - Returns: The public key as a DER-encoded `Data` blob ready to be sent to a server.
    @discardableResult
    func generateAndStoreKeyPair(config: KeyPairConfig) throws -> Data {
        // Delete any existing key with the same tag first
        try? deleteKeyPair(tag: config.tag)

        let privateKey: SecKey

        switch config.algorithm {
        case .ecSecureEnclave:
            privateKey = try generateSecureEnclaveKey(config: config)
        case .ecSoftware:
            privateKey = try generateECKey(config: config)
        case .rsa(let bits):
            privateKey = try generateRSAKey(bits: bits, config: config)
        }

        return try exportPublicKey(from: privateKey)
    }

    /// Retrieve the stored **private key** from the Keychain.
    func retrievePrivateKey(tag: String) throws -> SecKey {
        let query = baseKeychainQuery(tag: tag)
        var item: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &item)

        guard status == errSecSuccess, let key = item as! SecKey? else {
            throw KeyPairError.keyNotFound
        }
        return key
    }

    /// Export the **public key** (DER-encoded) for the stored key pair.
    func exportPublicKeyData(tag: String) throws -> Data {
        let privateKey = try retrievePrivateKey(tag: tag)
        return try exportPublicKey(from: privateKey)
    }

    /// Export the public key as a Base64-encoded PEM string.
    func exportPublicKeyPEM(tag: String) throws -> String {
        let der = try exportPublicKeyData(tag: tag)
        let b64 = der.base64EncodedString(options: [.lineLength64Characters])
        return "-----BEGIN PUBLIC KEY-----\n\(b64)\n-----END PUBLIC KEY-----"
    }

    /// Delete both keys associated with the given tag.
    func deleteKeyPair(tag: String) throws {
        let query = baseKeychainQuery(tag: tag)
        let status = SecItemDelete(query as CFDictionary)
        guard status == errSecSuccess || status == errSecItemNotFound else {
            throw KeyPairError.deletionFailed(status)
        }
    }

    // MARK: - Private: Key Generation

    private func generateSecureEnclaveKey(config: KeyPairConfig) throws -> SecKey {
        guard SecureEnclave.isAvailable else {
            throw KeyPairError.secureEnclaveUnavailable
        }

        var accessError: Unmanaged<CFError>?
        let flags: SecAccessControlCreateFlags = config.requiresBiometrics
            ? [.privateKeyUsage, .biometryCurrentSet]
            : [.privateKeyUsage]

        guard let access = SecAccessControlCreateWithFlags(
            kCFAllocatorDefault,
            kSecAttrAccessibleWhenUnlockedThisDeviceOnly,
            flags,
            &accessError
        ) else {
            throw accessError!.takeRetainedValue() as Error
        }

        let attributes: [String: Any] = [
            kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
            kSecAttrKeySizeInBits as String: 256,
            kSecAttrTokenID as String: kSecAttrTokenIDSecureEnclave,
            kSecPrivateKeyAttrs as String: [
                kSecAttrIsPermanent as String: true,
                kSecAttrApplicationTag as String: tagData(config.tag),
                kSecAttrAccessControl as String: access
            ]
        ]

        var error: Unmanaged<CFError>?
        guard let privateKey = SecKeyCreateRandomKey(attributes as CFDictionary, &error) else {
            throw error!.takeRetainedValue() as Error
        }
        return privateKey
    }

    private func generateECKey(config: KeyPairConfig) throws -> SecKey {
        var accessError: Unmanaged<CFError>?
        let flags: SecAccessControlCreateFlags = config.requiresBiometrics
            ? [.biometryCurrentSet]
            : []

        guard let access = SecAccessControlCreateWithFlags(
            kCFAllocatorDefault,
            kSecAttrAccessibleWhenUnlockedThisDeviceOnly,
            flags,
            &accessError
        ) else {
            throw accessError!.takeRetainedValue() as Error
        }

        let attributes: [String: Any] = [
            kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
            kSecAttrKeySizeInBits as String: 256,
            kSecPrivateKeyAttrs as String: [
                kSecAttrIsPermanent as String: true,
                kSecAttrApplicationTag as String: tagData(config.tag),
                kSecAttrAccessControl as String: access
            ]
        ]

        var error: Unmanaged<CFError>?
        guard let privateKey = SecKeyCreateRandomKey(attributes as CFDictionary, &error) else {
            throw error!.takeRetainedValue() as Error
        }
        return privateKey
    }

    private func generateRSAKey(bits: Int, config: KeyPairConfig) throws -> SecKey {
        var accessError: Unmanaged<CFError>?
        let flags: SecAccessControlCreateFlags = config.requiresBiometrics
            ? [.biometryCurrentSet]
            : []

        guard let access = SecAccessControlCreateWithFlags(
            kCFAllocatorDefault,
            kSecAttrAccessibleWhenUnlockedThisDeviceOnly,
            flags,
            &accessError
        ) else {
            throw accessError!.takeRetainedValue() as Error
        }

        let attributes: [String: Any] = [
            kSecAttrKeyType as String: kSecAttrKeyTypeRSA,
            kSecAttrKeySizeInBits as String: bits,
            kSecPrivateKeyAttrs as String: [
                kSecAttrIsPermanent as String: true,
                kSecAttrApplicationTag as String: tagData(config.tag),
                kSecAttrAccessControl as String: access
            ]
        ]

        var error: Unmanaged<CFError>?
        guard let privateKey = SecKeyCreateRandomKey(attributes as CFDictionary, &error) else {
            throw error!.takeRetainedValue() as Error
        }
        return privateKey
    }

    // MARK: - Private: Export

    private func exportPublicKey(from privateKey: SecKey) throws -> Data {
        guard let publicKey = SecKeyCopyPublicKey(privateKey) else {
            throw KeyPairError.publicKeyExportFailed(errSecParam)
        }
        var error: Unmanaged<CFError>?
        guard let data = SecKeyCopyExternalRepresentation(publicKey, &error) as Data? else {
            throw error!.takeRetainedValue() as Error
        }
        return data
    }

    // MARK: - Private: Helpers

    private func tagData(_ tag: String) -> Data {
        tag.data(using: .utf8)!
    }

    private func baseKeychainQuery(tag: String) -> [String: Any] {
        [
            kSecClass as String: kSecClassKey,
            kSecAttrKeyClass as String: kSecAttrKeyClassPrivate,
            kSecAttrApplicationTag as String: tagData(tag),
            kSecReturnRef as String: true
        ]
    }
}

// MARK: - Convenience: Sign & Verify

extension KeyPairManager {

    /// Sign arbitrary data with the stored private key.
    func sign(data: Data, tag: String) throws -> Data {
        let privateKey = try retrievePrivateKey(tag: tag)
        let algorithm = SecKeyAlgorithm.ecdsaSignatureMessageX962SHA256

        guard SecKeyIsAlgorithmSupported(privateKey, .sign, algorithm) else {
            throw KeyPairError.unsupportedAlgorithm
        }

        var error: Unmanaged<CFError>?
        guard let signature = SecKeyCreateSignature(
            privateKey,
            algorithm,
            data as CFData,
            &error
        ) as Data? else {
            throw error!.takeRetainedValue() as Error
        }
        return signature
    }

    /// Verify a signature using the public key derived from the stored private key.
    func verify(data: Data, signature: Data, tag: String) throws -> Bool {
        let privateKey = try retrievePrivateKey(tag: tag)
        guard let publicKey = SecKeyCopyPublicKey(privateKey) else {
            throw KeyPairError.keyNotFound
        }
        let algorithm = SecKeyAlgorithm.ecdsaSignatureMessageX962SHA256

        var error: Unmanaged<CFError>?
        let valid = SecKeyVerifySignature(
            publicKey,
            algorithm,
            data as CFData,
            signature as CFData,
            &error
        )
        return valid
    }
}

// MARK: - Demo / Usage

/// Call this from your AppDelegate, a ViewModel, or a SwiftUI .task modifier.
func runKeyPairDemo() {
    let manager = KeyPairManager()

    // --- 1. Secure Enclave (recommended on real devices) ---
    let seConfig = KeyPairManager.KeyPairConfig(
        tag: "com.myapp.userkey.se",
        algorithm: .ecSecureEnclave,
        requiresBiometrics: false
    )

    // --- 2. EC in software keychain (works on simulator too) ---
    let ecConfig = KeyPairManager.KeyPairConfig(
        tag: "com.myapp.userkey.ec",
        algorithm: .ecSoftware,
        requiresBiometrics: false
    )

    // --- 3. RSA 2048 ---
    let rsaConfig = KeyPairManager.KeyPairConfig(
        tag: "com.myapp.userkey.rsa",
        algorithm: .rsa(bits: 2048),
        requiresBiometrics: false
    )

    // Pick the appropriate config:
    let config = SecureEnclave.isAvailable ? seConfig : ecConfig

    do {
        // Generate & store
        let publicKeyData = try manager.generateAndStoreKeyPair(config: config)
        print("✅ Key pair generated. Public key (\(publicKeyData.count) bytes).")

        // Export PEM for upload to server
        let pem = try manager.exportPublicKeyPEM(tag: config.tag)
        print("📤 Public Key PEM:\n\(pem)")

        // Sign something
        let payload = "Hello, secure world!".data(using: .utf8)!
        let signature = try manager.sign(data: payload, tag: config.tag)
        print("✍️  Signature (\(signature.count) bytes): \(signature.base64EncodedString())")

        // Verify
        let valid = try manager.verify(data: payload, signature: signature, tag: config.tag)
        print("🔍 Signature valid: \(valid)")

        // Clean up (optional)
        // try manager.deleteKeyPair(tag: config.tag)

    } catch {
        print("❌ Error: \(error.localizedDescription)")
    }
}