Sesame-iOS/Sesame/KeyManagement.swift

import Foundation
import CryptoKit
import SwiftUI

final class KeyManagement {
    
    static let securityKeySize: SymmetricKeySize = .bits128
    
    enum KeyError: Error {
        /// Keys which are already in use can't be exported
        case exportAttemptOfUsedKeys
    }
    
    static var documentsDirectory: URL {
        let paths = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask)
        return paths[0]
    }
    
    private let keyFile = KeyManagement.documentsDirectory.appendingPathComponent("keys")
    
    let exportFile = KeyManagement.documentsDirectory.appendingPathComponent("export.cpp")
    
    private var keys: [(key: SymmetricKey, used: Bool)] {
        didSet {
            do {
                try saveKeys()
            } catch {
                print("Failed to save changed keys: \(error)")
            }
        }
    }
    
    var numberOfKeys: Int {
        keys.count
    }
    
    var hasUsedKeys: Bool {
        keys.contains { $0.used }
    }
    
    var hasUnusedKeys: Bool {
        unusedKeyCount > 0
    }
    
    var unusedKeyCount: Int {
        guard let id = nextKeyId else {
            return 0
        }
        return keys.count - id + 1
    }
    
    var usedKeyCount: Int {
        nextKeyId ?? keys.count
    }
    
    var lastKeyId: Int? {
        keys.lastIndex { $0.used }
    }
    
    var nextKeyId: Int? {
        let index = lastKeyId ?? -1 + 1
        guard index < keys.count else {
            return nil
        }
        return index
    }
    
    init() throws {
        guard FileManager.default.fileExists(atPath: keyFile.path) else {
            self.keys = []
            return
        }
        let content = try String(contentsOf: keyFile)
        self.keys = content.components(separatedBy: "\n")
            .enumerated().compactMap { (index, line) -> (SymmetricKey, Bool)? in
                let parts = line.components(separatedBy: ":")
                guard parts.count == 2 else {
                    return nil
                }
                let keyData = Data(base64Encoded: parts[0])!
                return (SymmetricKey(data: keyData), parts[1] != "0")
            }
        print("\(unusedKeyCount) / \(keys.count) keys remaining")
    }
    
    func useNextKey() -> (key: SymmetricKey, id: Int)? {
        guard let index = nextKeyId else {
            return nil
        }
        let key = keys[index].key
        keys[index].used = true
        return (key, index)
    }
    
    func regenerateKeys(count: Int = 100) throws {
        self.keys = Self.generateKeys(count: count)
            .map { ($0, false) }
        let keyString = keys.map { $0.key.codeString }.joined(separator: "\n")
        try keyString.write(to: exportFile, atomically: false, encoding: .utf8)
    }
    
    private func saveKeys() throws {
        let content = keys.map { key, used -> String in
            let keyString = key.withUnsafeBytes {
                return Data(Array($0)).base64EncodedString()
            }
            return keyString + ":" + (used ? "1" : "0")
        }.joined(separator: "\n")
        try content.write(to: keyFile, atomically: true, encoding: .utf8)
        print("Keys saved")
    }
    
    static func generateKeys(count: Int = 100) -> [SymmetricKey] {
        (0..<count).map { _ in
            SymmetricKey(size: securityKeySize)
        }
    }
}