Thijooree-iOS/Thijooree iOS/Storage/CacheEncryption.swift

import Foundation
import CryptoKit
import Security

enum CacheEncryptionError: Error, LocalizedError {
    case keyStorageFailed(OSStatus)
    case encryptionFailed
    case decryptionFailed
    case invalidData

    var errorDescription: String? {
        switch self {
        case .keyStorageFailed(let s): return "Keychain key storage failed (OSStatus \(s))"
        case .encryptionFailed:        return "AES-GCM encryption failed"
        case .decryptionFailed:        return "AES-GCM decryption failed"
        case .invalidData:             return "Data is not valid Base64 or UTF-8"
        }
    }
}

// AES-256-GCM encryption for all local caches.
// Mirrors Android's CacheEncryption (AES-256-GCM, AndroidKeyStore-backed key).
// The symmetric key is generated once and stored in Keychain, inaccessible when device is locked.
final class CacheEncryption {
    static let shared = CacheEncryption()

    private let keychainService = "sh.sar.thijooree.cachekey"
    private let keychainAccount = "aes256_master_key"
    private let lock = NSLock()
    private var cachedKey: SymmetricKey?

    private init() {}

    // MARK: - String convenience

    func encrypt(_ string: String) throws -> String {
        guard let data = string.data(using: .utf8) else {
            throw CacheEncryptionError.invalidData
        }
        return try encryptData(data)
    }

    func decryptString(_ base64: String) throws -> String {
        let data = try decryptData(base64)
        guard let string = String(data: data, encoding: .utf8) else {
            throw CacheEncryptionError.invalidData
        }
        return string
    }

    // MARK: - Data operations

    // Returns a Base64 string containing the AES-GCM sealed box (nonce + ciphertext + tag).
    func encryptData(_ plaintext: Data) throws -> String {
        let key = try loadOrCreateKey()
        guard let combined = (try? AES.GCM.seal(plaintext, using: key))?.combined else {
            throw CacheEncryptionError.encryptionFailed
        }
        return combined.base64EncodedString()
    }

    // Decrypts a Base64 AES-GCM sealed box back to the original data.
    func decryptData(_ base64: String) throws -> Data {
        guard let combined = Data(base64Encoded: base64) else {
            throw CacheEncryptionError.invalidData
        }
        let key = try loadOrCreateKey()
        do {
            return try AES.GCM.open(AES.GCM.SealedBox(combined: combined), using: key)
        } catch {
            throw CacheEncryptionError.decryptionFailed
        }
    }

    // MARK: - Key lifecycle

    // Clears the cached key and removes it from Keychain.
    // Call on full logout so cached data becomes unreadable.
    func purge() {
        lock.lock()
        defer { lock.unlock() }
        cachedKey = nil
        let query: [String: Any] = [
            kSecClass as String:       kSecClassGenericPassword,
            kSecAttrService as String: keychainService,
            kSecAttrAccount as String: keychainAccount
        ]
        SecItemDelete(query as CFDictionary)
    }

    // MARK: - Private

    private func loadOrCreateKey() throws -> SymmetricKey {
        lock.lock()
        defer { lock.unlock() }

        if let key = cachedKey { return key }

        if let key = loadKeyFromKeychain() {
            cachedKey = key
            return key
        }

        let newKey = SymmetricKey(size: .bits256)
        try saveKeyToKeychain(newKey)
        cachedKey = newKey
        return newKey
    }

    private func loadKeyFromKeychain() -> SymmetricKey? {
        let query: [String: Any] = [
            kSecClass as String:       kSecClassGenericPassword,
            kSecAttrService as String: keychainService,
            kSecAttrAccount as String: keychainAccount,
            kSecReturnData as String:  true,
            kSecMatchLimit as String:  kSecMatchLimitOne
        ]
        var result: AnyObject?
        guard SecItemCopyMatching(query as CFDictionary, &result) == errSecSuccess,
              let data = result as? Data, data.count == 32 else { return nil }
        return SymmetricKey(data: data)
    }

    private func saveKeyToKeychain(_ key: SymmetricKey) throws {
        let keyData = key.withUnsafeBytes { Data($0) }

        let searchQuery: [String: Any] = [
            kSecClass as String:       kSecClassGenericPassword,
            kSecAttrService as String: keychainService,
            kSecAttrAccount as String: keychainAccount
        ]
        let attributes: [String: Any] = [
            kSecValueData as String:    keyData,
            kSecAttrAccessible as String: kSecAttrAccessibleWhenUnlockedThisDeviceOnly
        ]

        let addQuery = searchQuery.merging(attributes) { _, new in new }
        let addStatus = SecItemAdd(addQuery as CFDictionary, nil)

        if addStatus == errSecDuplicateItem {
            let updateStatus = SecItemUpdate(searchQuery as CFDictionary, attributes as CFDictionary)
            guard updateStatus == errSecSuccess else {
                throw CacheEncryptionError.keyStorageFailed(updateStatus)
            }
        } else if addStatus != errSecSuccess {
            throw CacheEncryptionError.keyStorageFailed(addStatus)
        }
    }
}