gulliver / Sources / Gulliver / KeychainStorage.swift

//
//  File.swift
//  Gulliver
//
//  Created by Bailey Townsend on 1/20/26.
//

import Foundation
import CryptoKit

enum SecureStore: String {
    case state = "state"
    case dpop = "dpop"
}



public func getStateKeychainStore() -> KeychainStorage {    
    return KeychainStorage(namespace: "state")
}

public func getSessionKeychainStore() -> KeychainStorage {
    return KeychainStorage(namespace: "session")
}



public class KeychainStorage {
    let namespace: String
    private let accessGroup: String?
    
    private static var defaultAccessibility: CFString {
#if os(iOS)
        return kSecAttrAccessibleAfterFirstUnlock
#elseif os(macOS)
        return kSecAttrAccessibleAfterFirstUnlock
#endif
    }
    
    private static func platformSpecificAttributes() -> [String: Any] {
        var attributes: [String: Any] = [:]
        
#if os(macOS)
        // Disable iCloud sync for app-specific keychain items on macOS
        attributes[kSecAttrSynchronizable as String] = false
#endif
        
        return attributes
    }
    
    
    public init(namespace: String = "default", accessGroup: String? = nil) {
        self.namespace = namespace
        self.accessGroup = accessGroup
    }
    
    private func getKey(did: String, store: SecureStore ) -> String {
        "\(did):\(store)"
    }
    
    
    
    func store(key: String, value: Data, namespace: String) throws {
        let namespacedKey = "\(namespace):\(key)"
        
        var query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: namespacedKey,
            kSecValueData as String: value,
            kSecAttrAccessible as String: Self.defaultAccessibility,
        ]
        
        // Add platform-specific attributes
        query.merge(Self.platformSpecificAttributes()) { _, new in new }
        
        // Delete any existing item with the same key
        let deleteStatus = SecItemDelete(query as CFDictionary)
        if deleteStatus != errSecSuccess, deleteStatus != errSecItemNotFound {
            throw KeychainStorageError.deleteError(Int(deleteStatus))
        }
        
        // Add the new item to the keychain
        let status = SecItemAdd(query as CFDictionary, nil)
        if status == errSecDuplicateItem {
            throw KeychainStorageError.keyStoreError(Int(status))
        }
        guard status == errSecSuccess else {
            throw KeychainStorageError.keyStoreError(Int(status))
        }
    }
    
    func retrieve(key: String, namespace: String) throws -> Data {
        let namespacedKey = "\(namespace):\(key)"
        
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: namespacedKey,
            kSecReturnData as String: kCFBooleanTrue!,
            kSecMatchLimit as String: kSecMatchLimitOne,
        ]
        
        var item: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &item)
        
        if status == errSecItemNotFound {
            throw KeychainStorageError.retrieveError(Int(status))
        }
        
        guard status == errSecSuccess else {
            throw KeychainStorageError.retrieveError(Int(status))
        }
        guard let data = item as? Data else {
            throw KeychainStorageError.dataFormatError
        }
        
        return data
    }
    
    func delete(key: String, namespace: String) throws {
        let namespacedKey = "\(namespace):\(key)"
        
        
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrAccount as String: namespacedKey,
        ]
        
        let status = SecItemDelete(query as CFDictionary)
        
        if status != errSecSuccess, status != errSecItemNotFound {
            throw KeychainStorageError.deleteError(Int(status))
        }
    }
    
    func deleteAll(namespace: String) throws {
        // Handle generic passwords first
        let genericSuccess = try deleteGenericPasswords(withNamespacePrefix: namespace)
        
        // Then handle crypto keys
        let keysSuccess = try deleteCryptoKeys(withNamespacePrefix: namespace)
        
        guard genericSuccess, keysSuccess else {
            throw KeychainStorageError.deleteError(-1)
        }
    }
    
    private func deleteGenericPasswords(withNamespacePrefix namespace: String) throws -> Bool {
        // Query to get all generic passwords
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecMatchLimit as String: kSecMatchLimitAll,
            kSecReturnAttributes as String: true,
        ]
        
        var result: AnyObject?
        let status = SecItemCopyMatching(query as CFDictionary, &result)
        
        if status == errSecSuccess, let items = result as? [[String: Any]] {
            var allSucceeded = true
            var matchedCount = 0
            
            // Filter and delete items that match our namespace
            for item in items {
                if let account = item[kSecAttrAccount as String] as? String,
                   account.hasPrefix("\(namespace):")
                {
                    matchedCount += 1
                    
                    let deleteQuery: [String: Any] = [
                        kSecClass as String: kSecClassGenericPassword,
                        kSecAttrAccount as String: account,
                    ]
                    
                    let deleteStatus = SecItemDelete(deleteQuery as CFDictionary)
                    if deleteStatus != errSecSuccess {
                        allSucceeded = false
                    }
                }
            }
            
            return allSucceeded
        } else if status == errSecItemNotFound {
            
            return true
        } else {
            //TODO throw or just sliently fail?
            return false
        }
    }
    
    private func deleteCryptoKeys(withNamespacePrefix namespace: String) throws -> Bool {
        // Query to get all keys
        let query: [String: Any] = [
            kSecClass as String: kSecClassKey,
            kSecMatchLimit as String: kSecMatchLimitAll,
            kSecReturnAttributes as String: true,
        ]
        
        var result: AnyObject?
        let status = SecItemCopyMatching(query as CFDictionary, &result)
        
        if status == errSecSuccess, let items = result as? [[String: Any]] {
            var allSucceeded = true
            var matchedCount = 0
            
            // Filter and delete keys
            for item in items {
                // For keys, check the application tag
                if let tagData = item[kSecAttrApplicationTag as String] as? Data,
                   let tagString = String(data: tagData, encoding: .utf8),
                   tagString.hasPrefix("\(namespace).")
                {
                    matchedCount += 1
                    
                    let deleteQuery: [String: Any] = [
                        kSecClass as String: kSecClassKey,
                        kSecAttrApplicationTag as String: tagData,
                    ]
                    
                    let deleteStatus = SecItemDelete(deleteQuery as CFDictionary)
                    if deleteStatus != errSecSuccess {
                        allSucceeded = false
                    }
                }
            }
            
            return allSucceeded
        } else if status == errSecItemNotFound {
            return true
        } else {
            return false
        }
    }
    
    func storeDPoPKey(_ key: P256.Signing.PrivateKey, keyTag: String) throws {
        guard let tagData = keyTag.data(using: .utf8) else {
            throw KeychainStorageError.dataFormatError
        }
        
        var query: [String: Any] = [
            kSecClass as String: kSecClassKey,
            kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
            kSecAttrKeySizeInBits as String: 256,
            kSecAttrApplicationTag as String: tagData,
            kSecValueData as String: key.x963Representation,
            kSecAttrAccessible as String: Self.defaultAccessibility,
        ]
        
        // Delete any existing key first
        let deleteQuery: [String: Any] = [
            kSecClass as String: kSecClassKey,
            kSecAttrApplicationTag as String: tagData,
        ]
        
        let deleteStatus = SecItemDelete(deleteQuery as CFDictionary)
        
        // Add the new key
        let status = SecItemAdd(query as CFDictionary, nil)
        if status == errSecDuplicateItem {
            // Try update
            let updateAttributes: [String: Any] = [
                kSecValueData as String: key.x963Representation,
            ]
            let updateStatus = SecItemUpdate(deleteQuery as CFDictionary, updateAttributes as CFDictionary)
            guard updateStatus == errSecSuccess else {
                throw KeychainStorageError.keyStoreError(Int(updateStatus))
            }
        } else if status != errSecSuccess {
            throw KeychainStorageError.keyStoreError(Int(status))
        }
    }
    
    func retrieveDPoPKey(keyTag: String) throws -> P256.Signing.PrivateKey {
        guard let tagData = keyTag.data(using: .utf8) else {
            throw KeychainStorageError.dataFormatError
        }
        
        let query: [String: Any] = [
            kSecClass as String: kSecClassKey,
            kSecAttrApplicationTag as String: tagData,
            kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
            kSecReturnData as String: kCFBooleanTrue!,
            kSecMatchLimit as String: kSecMatchLimitOne,
        ]
        
        var item: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &item)
        
        guard status == errSecSuccess, let data = item as? Data else {
            throw KeychainStorageError.retrieveError(Int(status))
        }
                
        return  try P256.Signing.PrivateKey(x963Representation: data)
                    
    }
    
    func deleteDPoPKey(keyTag: String) throws {
        guard let tagData = keyTag.data(using: .utf8) else {
            throw KeychainStorageError.dataFormatError
        }
        
        let query: [String: Any] = [
            kSecClass as String: kSecClassKey,
            kSecAttrApplicationTag as String: tagData,
        ]
        
        let status = SecItemDelete(query as CFDictionary)
        if status != errSecSuccess, status != errSecItemNotFound {
            throw KeychainStorageError.deleteError(Int(status))
        }
    }
    
}



public enum KeychainStorageError: Error {
    case deleteError(Int)
    case keyStoreError(Int)
    case retrieveError(Int)
    case dataFormatError
    
}