NorthSec 2017 Badge Customization

nsec17_badge_client.swift

//
//  nsec17_badge_client.swift
//  BadgeClient
//
//  Created by Marc-Etienne M.Léveillé on 2017-05-16.
//  Copyright © 2017 NorthSec. All rights reserved.
//  License: BSD 2-clause
//

import Foundation
import CoreBluetooth

let kIdentityServiceUUID = CBUUID(string: "6AA99996-93B3-95BA-BC40-42684CDB36FB")

let kIdentityCharNameUUID      = CBUUID(string: "6AA96E6D-93B3-95BA-BC40-42684CDB36FB")
let kIdentityCharAvatarUUID    = CBUUID(string: "6AA96176-93B3-95BA-BC40-42684CDB36FB")
let kIdentityCharUnlockKeyUUID = CBUUID(string: "6AA96B79-93B3-95BA-BC40-42684CDB36FB")

func printUsageAndExit() {
    print("Usage: \(CommandLine.arguments.first!) --device NSEC1234 --key 5678 [--name NAME | --avatar BITMAP_FILE]")
    exit(EXIT_FAILURE)
}

var args: [String] = CommandLine.arguments.suffix(from: 1).reversed()

var device: String?  = nil
var unlockKey: String? = nil
var newName: String? = nil
var avatarFilename: String? = nil

while args.count >= 2 {
    switch(args.popLast()!) {
        case "--device":
            device = args.popLast()
        case "--key":
            unlockKey = args.popLast()
        case "--name":
            newName = args.popLast()
        case "--avatar":
            avatarFilename = args.popLast()
        default:
            printUsageAndExit()
    }
}

if(device == nil || unlockKey == nil || (newName == nil && avatarFilename == nil)) {
    printUsageAndExit()
}

enum UpdateState {
    case locked, sendName, sendAvatar, needRelock
}

var connectedDevice: CBPeripheral? = nil

class Identity: NSObject, CBCentralManagerDelegate, CBPeripheralDelegate {
    var charateristics: [CBUUID : CBCharacteristic] = [:]
    var state: UpdateState = UpdateState.locked

    func centralManagerDidUpdateState(_ central: CBCentralManager) {
        switch central.state {
        case CBCentralManagerState.poweredOn:
            print("Scanning...")
            central.scanForPeripherals(withServices: [kIdentityServiceUUID], options: nil)
        default:
            break
        }
    }

    func centralManager(_ central: CBCentralManager, didDiscover peripheral: CBPeripheral, advertisementData: [String : Any], rssi RSSI: NSNumber) {
        print("Found badge", peripheral.name ?? "<Unkonwn>")
        if peripheral.name == device {
            print("Connecting to", peripheral.name ?? "<Unkonwn>")
            connectedDevice = peripheral
            central.connect(peripheral, options: nil)
        }
    }

    func centralManager(_ central: CBCentralManager, didConnect peripheral: CBPeripheral) {
        print("Connected")
        peripheral.delegate = self
        peripheral.discoverServices([kIdentityServiceUUID])
    }

    func centralManager(_ central: CBCentralManager, didFailToConnect peripheral: CBPeripheral, error: Error?) {
        print("Error", error as Any)
    }

    func peripheral(_ peripheral: CBPeripheral, didDiscoverServices error: Error?) {
        peripheral.services?.filter({(x: CBService) -> Bool in return x.uuid.isEqual(kIdentityServiceUUID) }).forEach({ (s) in
            peripheral.discoverCharacteristics(nil, for: s)
        })
    }

    func peripheral(_ peripheral: CBPeripheral, didDiscoverCharacteristicsFor service: CBService, error: Error?) {
        service.characteristics?.forEach({ (c) in
            charateristics[c.uuid] = c
        })
        trySendData()
    }

    func trySendData() {
        switch state {
        case .locked:
            if let c = charateristics[kIdentityCharUnlockKeyUUID] {
                print("Unlocking badge")
                connectedDevice?.writeValue(unlockKey?.data(using: String.Encoding.utf8) ?? Data(),
                                      for:c,
                                      type: CBCharacteristicWriteType.withResponse)
            }
        case .sendName:
            if let c = charateristics[kIdentityCharNameUUID] {
                print("Updating name")
                connectedDevice?.writeValue(newName?.data(using: String.Encoding.utf8) ?? Data(),
                                            for:c,
                                            type: CBCharacteristicWriteType.withResponse)
            }
        case .sendAvatar:
            if let c = charateristics[kIdentityCharAvatarUUID] {
                print("Updating avatar")
                do {
                    try connectedDevice?.writeValue(Data(contentsOf: URL(fileURLWithPath: avatarFilename!)),
                                                for:c,
                                                type: CBCharacteristicWriteType.withResponse)
                }
                catch {
                    // oups
                }
            }
        case .needRelock:
            if let c = charateristics[kIdentityCharUnlockKeyUUID] {
                print("Relocking badge")
                connectedDevice?.writeValue(Data(), for:c, type: CBCharacteristicWriteType.withResponse)
            }
        }
    }

    func peripheral(_ peripheral: CBPeripheral, didWriteValueFor characteristic: CBCharacteristic, error: Error?) {
        switch characteristic.uuid {
        case kIdentityCharUnlockKeyUUID:
            if state == .needRelock {
                print("Relocked")
                exit(EXIT_SUCCESS)
            }
            else if newName != nil {
                state = .sendName
            }
            else {
                state = .sendAvatar
            }
        case kIdentityCharNameUUID:
            if avatarFilename != nil {
                state = .sendAvatar
            }
            else {
                state = .needRelock
            }
        case kIdentityCharAvatarUUID:
            state = .needRelock
        default:
            break
        }
        trySendData()
    }
}

var client = Identity()
var central = CBCentralManager(delegate: client, queue: nil)

RunLoop.main.run()