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)? /// 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 { 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") } }