Sesame-Server/Sources/App/KeyManagement.swift

import Foundation
import WebSocketKit
import Vapor

final class KeyManagement {
    
    /// The security parameter for the keys (in bits)
    private static let keySecurity = 128
    
    /// The size of the individual keys in bytes
    static let keySize = keySecurity / 8
    
    /// The seconds to wait for a response from the device
    static let deviceTimeout: Int64 = 20

    /// The connection to the device
    private var connection: WebSocket?
    
    private let deviceKey: String
    
    var deviceIsAuthenticated = false
    
    /// Indicator for device availability
    var deviceIsConnected: Bool {
        !(connection?.isClosed ?? true) && deviceIsAuthenticated
    }
    
    /// The id of the key which was sent to the device
    private var keyInTransit: (id: UInt16, promise: EventLoopPromise<KeyResult>)?
    
    /// The result transmitted by the device for the sent key
    var keyResult: KeyResult = .none
    
    init(deviceKey: String) {
        self.deviceKey = deviceKey
    }
    
    // MARK: API
    
    var deviceResponse: String {
        guard let keyId = keyInTransit else {
            return "No key"
        }
        return "\(keyId):\(keyResult.rawValue)"
    }
    
    var deviceStatus: String {
        deviceIsConnected ? "1" : "0"
    }
    
    func clearClientRequest() {
        keyInTransit = nil
        keyResult = .none
    }
    
    func sendKeyToDevice(_ key: Data, keyId: UInt16, on eventLoop: EventLoop) -> EventLoopFuture<KeyResult> {
        guard key.count == KeyManagement.keySize else {
            return eventLoop.makeSucceededFuture(.invalidPayloadSize)
        }
        guard let socket = connection, !socket.isClosed else {
            connection = nil
            return eventLoop.makeSucceededFuture(.deviceNotConnected)
        }
        let keyIdData = [UInt8(keyId >> 8), UInt8(keyId & 0xFF)]
        let promise = eventLoop.makePromise(of: KeyResult.self)
        keyInTransit = (keyId, promise)
        socket.send(keyIdData + key, promise: nil)
        eventLoop.scheduleTask(in: .seconds(Self.deviceTimeout)) { [weak self] in
            guard let (storedKeyId, promise) = self?.keyInTransit, storedKeyId == keyId else {
                return
            }
            self?.keyInTransit = nil
            promise.succeed(.deviceTimedOut)
        }
        return promise.futureResult
    }
    
    func authenticateDevice(psk: String) {
        guard psk == self.deviceKey else {
            print("Invalid device key")
            _ = connection?.close()
            deviceIsAuthenticated = false
            return
        }
        print("Device authenticated")
        deviceIsAuthenticated = true
    }

    func processDeviceResponse(_ data: ByteBuffer) {
        guard let (_, promise) = keyInTransit else {
            print("No key in transit for response from device \(data)")
            return
        }
        defer { keyInTransit = nil }
        guard data.readableBytes == 1 else {
            print("Unexpected number of bytes received from device")
            promise.succeed(.unexpectedSocketEvent)
            return
        }
        guard let rawValue = data.getBytes(at: 0, length: 1)?.first else {
            print("Unreadable data received from device")
            promise.succeed(.unexpectedSocketEvent)
            return
        }
        guard let response = KeyResult(rawValue: rawValue) else {
            print("Unknown response \(rawValue) received from device")
            promise.succeed(.unexpectedSocketEvent)
            return
        }
        guard keyInTransit != nil else {
            print("No key in transit for response \(response)")
            return
        }
        keyResult = response
    }
    
    func didCloseDeviceSocket() {
        deviceIsAuthenticated = false
        guard connection != nil else {
            return
        }
        connection = nil
        print("Socket closed")
    }
    
    func removeDeviceConnection() {
        _ = connection?.close()
        connection = nil
    }
    
    func createNewDeviceConnection(_ socket: WebSocket) {
        removeDeviceConnection()
        connection = socket
        deviceIsAuthenticated = false
        print("Socket connected")
    }
}