learning_ai_common_plat/packages/swift-platform-sdk/Sources/BLFieldEncrypt.swift

// ── Field Encryption ────────────────────────────────────────
// AES-256-GCM field-level encryption compatible with @bytelyst/field-encrypt (TypeScript).
// Produces and consumes the same EncryptedField JSON structure across all platforms.
// Uses Apple CryptoKit — available on iOS 13+, macOS 10.15+, watchOS 6+.

import CryptoKit
import Foundation

// MARK: - EncryptedField Model

/// Encrypted field structure — wire-compatible with @bytelyst/field-encrypt (TypeScript).
/// All byte arrays are hex-encoded strings for JSON serialization.
public struct BLEncryptedField: Codable, Sendable, Equatable {
    /// Sentinel — always `true` for encrypted fields.
    public let __encrypted: Bool
    /// Schema version (currently 1).
    public let v: Int
    /// Algorithm identifier.
    public let alg: String
    /// Ciphertext (hex-encoded).
    public let ct: String
    /// Initialization vector (hex-encoded, 12 bytes = 24 hex chars).
    public let iv: String
    /// GCM authentication tag (hex-encoded, 16 bytes = 32 hex chars).
    public let tag: String
    /// DEK identifier — identifies which key was used for encryption.
    public let dekId: String

    public init(ct: String, iv: String, tag: String, dekId: String) {
        self.__encrypted = true
        self.v = 1
        self.alg = "aes-256-gcm"
        self.ct = ct
        self.iv = iv
        self.tag = tag
        self.dekId = dekId
    }
}

// MARK: - BLFieldEncrypt

/// AES-256-GCM field-level encryption.
///
/// Produces `BLEncryptedField` objects that are wire-compatible with the
/// TypeScript `@bytelyst/field-encrypt` package. Backends and native clients
/// can encrypt/decrypt the same fields interchangeably.
///
/// Usage:
/// ```swift
/// let key = BLFieldEncrypt.generateKey()
/// let encrypted = try BLFieldEncrypt.encrypt("sensitive data", key: key, dekId: "dek_user1_notes")
/// let decrypted = try BLFieldEncrypt.decrypt(encrypted, key: key)
/// ```
public enum BLFieldEncrypt {

    /// AES-256-GCM key size in bytes.
    public static let keySize = 32

    /// GCM nonce (IV) size in bytes.
    private static let nonceSize = 12

    // MARK: - Encrypt

    /// Encrypt a plaintext string with AES-256-GCM.
    ///
    /// - Parameters:
    ///   - plaintext: UTF-8 string to encrypt.
    ///   - key: 32-byte symmetric key.
    ///   - dekId: DEK identifier stored in the output for key lookup on decrypt.
    ///   - aad: Optional additional authenticated data (e.g., "userId:context").
    /// - Returns: `BLEncryptedField` with hex-encoded ciphertext, IV, and tag.
    /// - Throws: `BLFieldEncryptError` if key size is wrong or encryption fails.
    public static func encrypt(
        _ plaintext: String,
        key: SymmetricKey,
        dekId: String,
        aad: String? = nil
    ) throws -> BLEncryptedField {
        guard key.bitCount == keySize * 8 else {
            throw BLFieldEncryptError.invalidKeySize(expected: keySize, actual: key.bitCount / 8)
        }

        let plaintextData = Data(plaintext.utf8)
        let nonce = AES.GCM.Nonce()

        let sealedBox: AES.GCM.SealedBox
        if let aad = aad {
            sealedBox = try AES.GCM.seal(
                plaintextData,
                using: key,
                nonce: nonce,
                authenticating: Data(aad.utf8)
            )
        } else {
            sealedBox = try AES.GCM.seal(plaintextData, using: key, nonce: nonce)
        }

        return BLEncryptedField(
            ct: sealedBox.ciphertext.hexString,
            iv: Data(sealedBox.nonce).hexString,
            tag: sealedBox.tag.hexString,
            dekId: dekId
        )
    }

    // MARK: - Decrypt

    /// Decrypt a `BLEncryptedField` back to plaintext.
    ///
    /// - Parameters:
    ///   - field: Encrypted field object (from Cosmos DB or API response).
    ///   - key: 32-byte symmetric key (must match the key used to encrypt).
    ///   - aad: Optional AAD (must match the AAD used during encryption).
    /// - Returns: Decrypted UTF-8 string.
    /// - Throws: `BLFieldEncryptError` if decryption or authentication fails.
    public static func decrypt(
        _ field: BLEncryptedField,
        key: SymmetricKey,
        aad: String? = nil
    ) throws -> String {
        guard key.bitCount == keySize * 8 else {
            throw BLFieldEncryptError.invalidKeySize(expected: keySize, actual: key.bitCount / 8)
        }

        guard let ctData = Data(hexString: field.ct),
              let ivData = Data(hexString: field.iv),
              let tagData = Data(hexString: field.tag) else {
            throw BLFieldEncryptError.invalidHexEncoding
        }

        let nonce = try AES.GCM.Nonce(data: ivData)
        let sealedBox = try AES.GCM.SealedBox(nonce: nonce, ciphertext: ctData, tag: tagData)

        let decryptedData: Data
        if let aad = aad {
            decryptedData = try AES.GCM.open(sealedBox, using: key, authenticating: Data(aad.utf8))
        } else {
            decryptedData = try AES.GCM.open(sealedBox, using: key)
        }

        guard let plaintext = String(data: decryptedData, encoding: .utf8) else {
            throw BLFieldEncryptError.utf8DecodingFailed
        }

        return plaintext
    }

    // MARK: - Key Generation

    /// Generate a random 32-byte AES-256 symmetric key.
    public static func generateKey() -> SymmetricKey {
        SymmetricKey(size: .bits256)
    }

    /// Create a `SymmetricKey` from a hex-encoded string (64 hex chars = 32 bytes).
    public static func keyFromHex(_ hex: String) throws -> SymmetricKey {
        guard let data = Data(hexString: hex), data.count == keySize else {
            throw BLFieldEncryptError.invalidKeySize(expected: keySize, actual: hex.count / 2)
        }
        return SymmetricKey(data: data)
    }

    // MARK: - Keychain Key Derivation

    /// Get or create a persistent encryption key in the Keychain.
    ///
    /// On first call, generates a random 32-byte AES-256 key and stores it
    /// as a hex string in the Keychain. On subsequent calls, loads the
    /// existing key. This provides a stable per-device DEK for client-side
    /// encryption without requiring the backend to provision keys.
    ///
    /// - Parameters:
    ///   - service: Keychain service identifier (typically the app's bundle ID).
    ///   - account: Keychain account key (e.g., `"field_encrypt_dek"`).
    /// - Returns: A 32-byte `SymmetricKey` backed by Keychain storage.
    /// - Throws: `BLFieldEncryptError.invalidHexEncoding` if stored key is corrupt.
    public static func getOrCreateKey(service: String, account: String = "field_encrypt_dek") throws -> SymmetricKey {
        if let existingHex = BLKeychain.read(service: service, key: account) {
            return try keyFromHex(existingHex)
        }

        let newKey = generateKey()
        let hex = newKey.withUnsafeBytes { Data($0).hexString }
        BLKeychain.save(service: service, key: account, value: hex)
        return newKey
    }

    /// Load an existing encryption key from the Keychain without creating one.
    ///
    /// - Parameters:
    ///   - service: Keychain service identifier.
    ///   - account: Keychain account key (e.g., `"field_encrypt_dek"`).
    /// - Returns: The stored `SymmetricKey`, or `nil` if none exists.
    public static func loadKey(service: String, account: String = "field_encrypt_dek") -> SymmetricKey? {
        guard let hex = BLKeychain.read(service: service, key: account) else { return nil }
        return try? keyFromHex(hex)
    }

    /// Delete the stored encryption key from the Keychain.
    ///
    /// - Parameters:
    ///   - service: Keychain service identifier.
    ///   - account: Keychain account key.
    /// - Returns: `true` if the key was deleted or didn't exist.
    @discardableResult
    public static func deleteKey(service: String, account: String = "field_encrypt_dek") -> Bool {
        BLKeychain.delete(service: service, key: account)
    }

    // MARK: - Type Guard

    /// Check if a JSON value represents an encrypted field.
    /// Compatible with the TypeScript `isEncryptedField()` type guard.
    public static func isEncrypted(_ value: Any?) -> Bool {
        if let dict = value as? [String: Any],
           let sentinel = dict["__encrypted"] as? Bool,
           sentinel == true,
           dict["v"] != nil,
           dict["alg"] != nil,
           dict["ct"] != nil,
           dict["iv"] != nil,
           dict["tag"] != nil,
           dict["dekId"] != nil {
            return true
        }
        return false
    }

    /// Check if a `Codable` value is a `BLEncryptedField`.
    public static func isEncrypted(_ field: BLEncryptedField?) -> Bool {
        field?.__encrypted == true
    }
}

// MARK: - Errors

public enum BLFieldEncryptError: LocalizedError {
    case invalidKeySize(expected: Int, actual: Int)
    case invalidHexEncoding
    case utf8DecodingFailed

    public var errorDescription: String? {
        switch self {
        case .invalidKeySize(let expected, let actual):
            return "AES-256-GCM requires a \(expected)-byte key, got \(actual)"
        case .invalidHexEncoding:
            return "Failed to decode hex-encoded field data"
        case .utf8DecodingFailed:
            return "Decrypted data is not valid UTF-8"
        }
    }
}

// MARK: - Hex Helpers

extension Data {
    /// Hex-encode data to a lowercase string.
    var hexString: String {
        map { String(format: "%02x", $0) }.joined()
    }

    /// Initialize Data from a hex-encoded string.
    init?(hexString: String) {
        let hex = hexString.lowercased()
        guard hex.count.isMultiple(of: 2) else { return nil }

        var data = Data(capacity: hex.count / 2)
        var index = hex.startIndex
        while index < hex.endIndex {
            let nextIndex = hex.index(index, offsetBy: 2)
            guard let byte = UInt8(hex[index..<nextIndex], radix: 16) else { return nil }
            data.append(byte)
            index = nextIndex
        }
        self = data
    }
}