Sesame-iOS/Sesame/Common/KeyManagement.swift

import Foundation
import CryptoKit
import SwiftUI

struct KeySet {

    let remote: SymmetricKey

    let device: SymmetricKey

    let server: Data
}

extension KeyManagement {

    enum KeyType: String, Identifiable, CaseIterable {

        case deviceKey = "sesame-device"
        case remoteKey = "sesame-remote"
        case authToken = "sesame-remote-auth"

        var id: String {
            rawValue
        }

        var displayName: String {
            switch self {
            case .deviceKey:
                return "Unlock Key"
            case .remoteKey:
                return "Response Key"
            case .authToken:
                return "Server Token"
            }
        }

        var keyLength: SymmetricKeySize {
            .bits256
        }

        var usesHashing: Bool {
            switch self {
            case .authToken:
                return true
            default:
                return false
            }
        }
    }
}

extension KeyManagement.KeyType: CustomStringConvertible {

    var description: String {
        displayName
    }
}

private struct KeyChain {

    private let keyType = kSecAttrKeyTypeEC

    private let keyClass = kSecAttrKeyClassSymmetric

    private let domain: String

    init(domain: String) {
        self.domain = domain
    }

    private func baseQuery(for type: KeyManagement.KeyType) -> [String : Any] {
        [kSecClass as String: kSecClassInternetPassword,
         kSecAttrAccount as String: type.rawValue,
         kSecAttrServer as String: domain]
    }

    func save(_ type: KeyManagement.KeyType, _ key: SymmetricKey) {
        var query = baseQuery(for: type)
        query[kSecValueData as String] = key.data
        let status = SecItemAdd(query as CFDictionary, nil)
        guard status == errSecSuccess else {
            print("Failed to store \(type): \(status)")
            return
        }
        print("\(type) saved to keychain")
    }

    func load(_ type: KeyManagement.KeyType) -> SymmetricKey? {
        var query = baseQuery(for: type)
        query[kSecReturnData as String] = kCFBooleanTrue

        var item: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &item)
        guard status != -25300 else {
            return nil
        }
        guard status == errSecSuccess else {
            print("Failed to get \(type): \(status)")
            return nil
        }
        let key = item as! CFData
        return SymmetricKey(data: key as Data)
    }

    func delete(_ type: KeyManagement.KeyType) {
        let status = SecItemDelete(baseQuery(for: type) as CFDictionary)
        guard status == errSecSuccess || status == errSecItemNotFound else {
            print("Failed to remove \(type): \(status)")
            return
        }
        print("\(type) removed from keychain")
    }

    func has(_ type: KeyManagement.KeyType) -> Bool {
        load(type) != nil
    }
}

final class KeyManagement: ObservableObject {

    private let keyChain: KeyChain

    @Published
    private(set) var hasRemoteKey = false

    @Published
    private(set) var hasDeviceKey = false

    @Published
    private(set) var hasAuthToken = false

    @Published
    private(set) var hasAllKeys = false

    init() {
        self.keyChain = KeyChain(domain: "christophhagen.de")
        updateKeyStates()
    }

    func has(_ type: KeyType) -> Bool {
        switch type {
        case .deviceKey:
            return hasDeviceKey
        case .remoteKey:
            return hasRemoteKey
        case .authToken:
            return hasAuthToken
        }
    }

    func getAllKeys() -> KeySet? {
        guard let remoteKey = get(.remoteKey),
              let token = get(.authToken)?.data,
              let deviceKey = get(.deviceKey) else {
            return nil
        }
        return .init(remote: remoteKey, device: deviceKey, server: token)
    }

    func get(_ type: KeyType) -> SymmetricKey? {
        keyChain.load(type)
    }

    func delete(_ type: KeyType) {
        keyChain.delete(type)
        updateKeyStates()
    }

    func generate(_ type: KeyType) {
        let key = SymmetricKey(size: type.keyLength)
        save(type, key: key)
    }
    
    func save(_ type: KeyType, data: Data) {
        let key = SymmetricKey(data: data)
        save(type, key: key)
    }
    
    private func save(_ type: KeyType, key: SymmetricKey) {
        if keyChain.has(type) {
            keyChain.delete(type)
        }
        keyChain.save(type, key)
        updateKeyStates()
    }

    private func updateKeyStates() {
        self.hasRemoteKey = keyChain.has(.remoteKey)
        self.hasDeviceKey = keyChain.has(.deviceKey)
        self.hasAuthToken = keyChain.has(.authToken)
        self.hasAllKeys = hasRemoteKey && hasDeviceKey && hasAuthToken
    }
}