Sesame-iOS/Sesame/Common/RequestCoordinator.swift

import Foundation
import SwiftUI
import SwiftData

final class RequestCoordinator: ObservableObject {
    
    @Published
    var state: MessageResult = .noKeyAvailable
    
    @Published
    var pendingRequests: [PendingOperation] = []
    
    @Published
    var activeRequest: PendingOperation?
    
    @Published
    var keyManager = KeyManagement()
    
    @Published
    var isPerformingRequest: Bool = false
    
    @AppStorage("server")
    var serverPath: String = "https://christophhagen.de/sesame/"
    
    @AppStorage("localIP")
    var localAddress: String = "192.168.178.104/"
    
    @AppStorage("connectionType")
    var connectionType: ConnectionStrategy = .remoteFirst
    
    private let modelContext: ModelContext
    
    private let client = Client()
    
    private var timer: Timer?
    
    init(modelContext: ModelContext) {
        self.modelContext = modelContext
        if keyManager.hasAllKeys {
            self.state = .notChecked
        }
    }
    
    func checkConnection(using route: TransmissionType? = nil) {
        guard !isPerformingRequest else {
            return
        }
        isPerformingRequest = true
        Task {
            let route = route ?? connectionType.transmissionTypes.first!
            let (finalResult, _) = await performChallenge(route: route)
            DispatchQueue.main.async {
                self.state = finalResult.result
                self.isPerformingRequest = false
            }
            print("Finished connection test: \(finalResult)")
            scheduleReturnToReadyState()
        }
    }
    
    func startUnlock(quitAfterSuccess: Bool = false) {
        guard !isPerformingRequest else {
            return
        }
        isPerformingRequest = true
        Task {
            let finalResult = await performFullChallengeResponse()
            
            DispatchQueue.main.async {
                self.state = finalResult
                self.isPerformingRequest = false
            }
            if finalResult == .unlocked, quitAfterSuccess {
                DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(3)) {
                    exit(EXIT_SUCCESS)
                }
            }
            scheduleReturnToReadyState()
        }
    }
    
    private func performFullChallengeResponse() async -> MessageResult {
        let transmissionTypes = connectionType.transmissionTypes
        for route in transmissionTypes.dropLast() {
            if await performUnlockAndSaveItem(route: route) == .unlocked {
                return .unlocked
            }
        }
        guard let route = transmissionTypes.last else {
            // No transmission types at all
            return keyManager.hasAllKeys ? .notChecked : .noKeyAvailable
        }
        return await performUnlockAndSaveItem(route: route)
    }
    
    private func performUnlockAndSaveItem(route: TransmissionType) async -> MessageResult {
        let startTime = Date.now
        let result = await performFullChallengeResponse(route: route)
        let endTime = Date.now
        let roundTripTime = endTime.timeIntervalSince(startTime)
        print("Unlock took \(Int(roundTripTime * 1000)) ms (\(result.result))")
        let item = HistoryItem(message: result, startDate: startTime, route: route, finishDate: endTime)
        modelContext.insert(item)
        return result.result
    }
    
    private func performFullChallengeResponse(route: TransmissionType) async -> Message {
        let (challengeResponse, challenge) = await performChallenge(route: route)
        guard let challenge else {
            return challengeResponse
        }
        
        let (unlockResponse, secondaryChallenge) = await performUnlock(with: challenge.message, route: route)
        guard let secondaryChallenge else {
            return unlockResponse
        }
        let (secondUnlockResponse, _) = await performUnlock(with: secondaryChallenge.message, route: route)
        return secondUnlockResponse
    }
    
    private func performChallenge(route: TransmissionType) async -> ServerResponse {
        let initialMessage = Message.initial()
        let (result, challenge) = await send(initialMessage, route: route)
        guard let message = challenge?.message else {
            return (result, nil)
        }
        // Can't get here without the message being accepted
        guard message.messageType == .challenge else {
            print("Invalid message type for challenge: \(message)")
            return (result.with(result: .invalidMessageTypeFromDevice), nil)
        }
        return (result.with(result: .deviceAvailable), challenge)
    }
    
    private func performUnlock(with challenge: Message, route: TransmissionType) async -> ServerResponse {
        let request = challenge.requestMessage()
        let (unlockState, responseData) = await send(request, route: route)
        
        guard let response = responseData?.message else {
            return (unlockState, nil)
        }
        switch response.messageType {
        case .initial, .request:
            print("Invalid message type for response: \(response)")
            return (response.with(result: .invalidMessageTypeFromDevice), nil)
        case .challenge:
            // New challenge received, challenge was expired
            return (unlockState, responseData)
        case .response:
            break
        }
        
        guard response.serverChallenge == request.serverChallenge else {
            print("Invalid server challenge for unlock: \(response)")
            return (response.with(result: .invalidServerChallengeFromDevice), nil)
        }
        return (response.with(result: .unlocked), nil)
    }
    
    private func url(for route: TransmissionType) -> String {
        switch route {
        case .throughServer:
            return serverPath
        case .overLocalWifi:
            return localAddress
        }
    }
    
    private func send(_ message: Message, route: TransmissionType) async -> ServerResponse {
        guard let keys = keyManager.getAllKeys() else {
            return (message.with(result: .noKeyAvailable), nil)
        }
        let url = url(for: route)
        return await client.send(message, to: url, through: route, using: keys)
    }
    
    func resetState() {
        let hasKeys = keyManager.hasAllKeys
        DispatchQueue.main.async {
            self.state = hasKeys ? .notChecked : .noKeyAvailable
        }
    }

    func scheduleReturnToReadyState() {
        timer?.invalidate()
        
        DispatchQueue.main.async {
            self.timer = Timer.scheduledTimer(withTimeInterval: 3, repeats: false) { [weak self] timer in
                defer { timer.invalidate() }
                guard let self else {
                    return
                }
                self.resetState()
                self.timer = nil
            }
        }
    }
}