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 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: UInt16?
    
    /// 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) -> KeyPostResponse {
        guard key.count == KeyManagement.keySize else {
            return .invalidKeySize
        }
        guard let socket = connection, !socket.isClosed else {
            connection = nil
            return .deviceNotConnected
        }
        let keyIdData = [UInt8(keyId >> 8), UInt8(keyId & 0xFF)]
        keyInTransit = keyId
        socket.send(keyIdData + key, promise: nil)
        return .success
    }
    
    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 data.readableBytes == 1 else {
            print("Unexpected number of bytes received from device")
            keyInTransit = nil
            keyResult = .unexpectedSocketEvent
            return
        }
        guard let rawValue = data.getBytes(at: 0, length: 1)?.first else {
            print("Unreadable data received from device")
            keyInTransit = nil
            keyResult = .unexpectedSocketEvent
            return
        }
        guard let response = KeyResult(rawValue: rawValue) else {
            print("Unknown response \(rawValue) received from device")
            keyInTransit = nil
            keyResult = .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")
    }
}