/*
 *  Copyright © 2016, 2017, Oracle and/or its affiliates. All rights reserved.
 *  This software is dual-licensed to you under the MIT License (MIT) and
 *  the Universal Permissive License (UPL). See the LICENSE file in the root
 *  directory for license terms. You may choose either license, or both.
 */

import UIKit
import Foundation
import CocoaAsyncSocket

/**
 * The Bootstrapper
 * joins the UDP multicast group at the predefined address
 * *MULTICAST_ADDRESS* and port *UDP_PORT* to wait for a message
 * from the Network Provisioner for discovery or provisioning. If the message
 * received is for discovery *DISCOVER_REQUEST*, the Bootstrapper sends
 * back the client information. By default, the information is the client's MAC
 * address. This information will be displayed by the Network Provisioner to the
 * operator to select the target client. The information is sent as a key value
 * pair, where key="MAC" and value is the MAC address. If different information
 * is desired from the client, the *getDeviceInfo()* method should be
 * modified to return the desired data.
 *
 * If the message received is for provisioning *PROVISION_REQUEST*, the
 * Bootstrapper saves and verifies the provisioning data.
 * The result of the provisioning is sent back to the Network
 * Provisioner.
 *
 * The provisioning file should be in the unified provisioner format so that
 * the provisioning data is sent in encrypted form.
 *
 * Bootstrapper.run must be called before any other methods are called 
 * otherwise behavior is undefied.
 *
 * Users of Bootstrapper call the `run` method with a callback that is invoked
 * with the Bootstrapper status when there a response from the Network Provisioner
 * after a request for the provisioning file. The status may indicate "bad format"
 * or "an exception occurred" - 0x01, "download file does not exist" - 0x03, 
 * "file downloaded" 0x02.
 *
 * Once the file is downloaded and the application validates its contents
 * it must call `isProvisioned` to complete the communication with the
 * Network Provisioner.
 *
 * The application should call `stop()` to close the socket and free resources
 * when it done using the Bootstrapper.
 */

public class Bootstrapper {
    
    /**
     * The port to listen on for messages from the Network Provisioner.
     */
    static let UDP_PORT: UInt16 = 4456
    
    /**
     * The address to which the Network Provisioner sends multicast messages.
     */
    static let MULTICAST_ADDRESS: String = "238.163.7.96"
    
    /**
     * Message type: Network Provisioner to Bootstrapper- request for client identification information.
     */
    static let DISCOVER_REQUEST: UInt8 = 0x01
    
    /**
     * Message type: Bootstrapper to Network Provisioner- client identification information.
     */
    static let DISCOVER_RESPONSE: UInt8 = 0x02
    
    /**
     * Message type: Network Provisioner to Bootstrapper- provisioning information
     */
    static let PROVISION_REQUEST: UInt8 = 0x03
    
    /**
     * Message type: Bootstrapper to Network Provisioner- response provisioning status.
     */
    static let PROVISION_RESPONSE: UInt8 = 0x04
    
    /**
     * Status: Successful result
     */
    static let SUCCESS: UInt8 = 0x00
    
    /**
     * Status: Failed result
     */
    static let FAILURE: UInt8 = 0x01
    
    /**
     * Status: Unknown result
     */
    static let UNKNOWN: UInt8 = 0x02
    
    /**
     * Bootstrapper status
     */
    public var status: UInt8 = 0x00
    
    var START_PROVISION: Bool = false
    
    // Set only if run is called
    var inSocket : InSocket!
    
    var taStore : String!
    var taStorePassword: String!
    
    /**
     * The address from which the Network Provisioner sends messages.
     */
    var UDP_CLIENT_ADDRESS: String = ""
    
    /**
     * The port from which the Network Provisioner sends messages.
     */
    var UDP_CLIENT_PORT: UInt16 = 0
    
    /**
     * Bootstrapper calls this handler when it has received a response from a
     * request. It is not called for all Bootstrapper status changes. It is
     * intended to indicate that the file has been downloaded successfully or 
     * unsuccessfully.
     */
    private var provisionFileCallback: ((UInt8) -> ())? = nil
    private func runProvisionFileCallback() {
        if let provisionFileCallback = self.provisionFileCallback {
            provisionFileCallback(status)
        }
    }
    /*
    public init(taStore: String,
                provisionFileCallback: ((UInt8) -> ())? = nil) {
        self.taStore = taStore
        if inSocket != nil {
            inSocket.socket.close()
        }
        if let provisionFileCallback = provisionFileCallback {
            self.provisionFileCallback = provisionFileCallback
        }
        inSocket = InSocket(bootstrapper:self)
    }
    */
    
    public init(taStore: String, taStorePassword: String,
                provisionFileCallback: ((UInt8) -> ())? = nil) {
        self.taStore = taStore
        self.taStorePassword = taStorePassword
        if inSocket != nil {
            inSocket.socket.close()
        }
        if let provisionFileCallback = provisionFileCallback {
            self.provisionFileCallback = provisionFileCallback
        }
        inSocket = InSocket(bootstrapper:self)
    }
    
    public func stop() {
        DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(5)) {
            if self.inSocket != nil {
                self.inSocket.socket.close()
                self.inSocket = nil
            }
        }
    }
    
    /**
     * Sends the discovery response message.
     */
    func handleDiscoverRequest(client: GCDAsyncUdpSocket) -> [UInt8] {
        // client identification info
        let clientInfo: DeviceInfo = getDeviceInfo()
        var sendData: [UInt8] = []
        
        sendData.append(Bootstrapper.DISCOVER_RESPONSE)
        var offset = 1
        
        offset += encodeLengthValue(value: clientInfo.key, buffer: &sendData, offset: offset)
        _ = offset + encodeLengthValue(value: clientInfo.value, buffer: &sendData, offset: offset)
    
        return sendData
    }
    
    /**
     * Encode a string into a length value pair (LV).
     */
    func encodeLengthValue(value: String, buffer: inout [UInt8], offset: Int) -> Int {
    
        let bytes = Array(value.utf8)
    
        if bytes.count > 255 {
            print("Value encodes to over 255 bytes.")
            return 0
        }
        
        buffer.append(UInt8(bytes.count))
        buffer.append(contentsOf: bytes)
        
        return bytes.count + 1
    }
    
    
    private func getDeviceInfo() -> DeviceInfo {
        let mac = UIDevice.current.identifierForVendor?.uuidString
        return DeviceInfo(key: "MAC", value: mac!)
    }
    
    /**
     * Completes the provisioning of the device.
     */
    func handleProvisionRequest(request: [UInt8], offset: Int, length: Int) {
        var response: [UInt8]!
        var sendData: Foundation.Data!

        // Check for minimum possible valid length
        if length < 5 {
            print("\tProvisioning was unsuccessful, format was not correct.\n")
            response = provisionResponse(status: Bootstrapper.FAILURE)
            sendData = Data(bytes: response)
            inSocket.socket.send(sendData, toHost: UDP_CLIENT_ADDRESS,
                                 port: UDP_CLIENT_PORT,
                                 withTimeout: -1, tag: 0)
            status = 0x01
            runProvisionFileCallback()
            return
        }
 
        do {
            let _ = try self.updateStore(taStore: self.taStore, buff: request,
                                       offset: offset)

            if self.isProvisioned() {
                print("\tSuccessfully provisioned.\n")
                response = self.provisionResponse(status: Bootstrapper.SUCCESS)
                sendData = Data(bytes: response)
                self.inSocket.socket.send(sendData, toHost: self.UDP_CLIENT_ADDRESS,
                                     port: self.UDP_CLIENT_PORT, withTimeout: -1, tag: 0)
                self.status = 0x02
            } else {
                print("\tProvisioning was unsuccessful.\n")
                response = self.provisionResponse(status: Bootstrapper.FAILURE)
                sendData = Data(bytes: response)
                self.inSocket.socket.send(sendData, toHost: self.UDP_CLIENT_ADDRESS,
                                     port: self.UDP_CLIENT_PORT, withTimeout: -1, tag: 0)
                self.status = 0x03
            }
            
        } catch {
            print("Error \(error)")
            response = self.provisionResponse(status: Bootstrapper.FAILURE)
            sendData = Data(bytes: response)
            self.inSocket.socket.send(sendData, toHost: self.UDP_CLIENT_ADDRESS,
                                 port: self.UDP_CLIENT_PORT, withTimeout: -1, tag: 0)
            self.status = 0x01
        }
        self.runProvisionFileCallback()
    }
    
    /**
     * Sends the result status response message.
     */
    private func provisionResponse(status: UInt8) -> [UInt8] {
        var response: [UInt8] = []
        response.append(Bootstrapper.PROVISION_RESPONSE)
        response.append(status)
        return response
    }
    
    /**
     * Completes provisioning.
     */
    @discardableResult
    private func updateStore(taStore: String, buff: [UInt8], offset: Int) throws -> String {
        
        var length:UInt = UInt(buff.count)
        length -= UInt(offset)
        var newoffset:Int = offset
        
        // Need to move past the request length
        var reqlen:UInt = UInt(buff[newoffset] & 0xFF) << 8
        newoffset += 1
        length -= 1
        reqlen += (UInt)(buff[newoffset] & 0xFF)
        newoffset += 1
        length -= 1
        
        if length != reqlen {
            throw getNetworkClientError(error: NSLocalizedString(
               "Provisioning information too short", comment: ""))
            
        }
        
        let subbuff = buff[newoffset...Int(reqlen)]
        
        let data = Data(bytes: subbuff)
        
        FileManager.default.createFile(atPath: taStore, contents: data)
        if FileManager.default.fileExists(atPath: taStore) {
            return taStore
        } else {
            return ""
        }
    }
    /**
     * Check trusted asset store to determine if asset is already provisioned
     */
    private func isProvisioned() -> Bool {
        return FileManager.default.fileExists(atPath: taStore) &&
            validateTam(path: taStore, password: taStorePassword)
    }
    
    /*
    // Called by the application to confirm that the provisioning file
    // was validated
    public func isProvisioned() {
        var provisionResponse: [UInt8]!
        var sendData: Foundation.Data!
        print("\tSuccessfully provisioned.\n")
        provisionResponse = self.provisionResponse(status: Bootstrapper.SUCCESS)
        sendData = Data(bytes: provisionResponse)
        inSocket.socket.send(sendData, toHost: self.UDP_CLIENT_ADDRESS,
                            port: self.UDP_CLIENT_PORT,
                            withTimeout: -1, tag: 0)
        // HACK: Don't call the callback for this since the callback is only
        // until file is downloaded
        status = 0x02
    }
    */
    // Called by the application to indicate provisioning failed
    public func cancel() {
        var provisionResponse: [UInt8]!
        var sendData: Foundation.Data!
        print("\tProvisioning was unsuccessful.\n")
        provisionResponse = self.provisionResponse(status: Bootstrapper.FAILURE)
        sendData = Data(bytes: provisionResponse)
        inSocket.socket.send(sendData, toHost: self.UDP_CLIENT_ADDRESS,
                             port: self.UDP_CLIENT_PORT,
                             withTimeout: -1, tag: 0)
        // HACK: Don't call the callback for this since the callback is only
        // until file is downloaded
        status = 0x01
    }
}

/**
 * Device information to hold the identifying key and value.
 */
private class DeviceInfo {
    var key: String!
    var value: String!
    
    init(key: String, value: String) {
        self.key = key
        self.value = value
    }
}

class InSocket: NSObject, GCDAsyncUdpSocketDelegate {
    
    @objc let IP = Bootstrapper.MULTICAST_ADDRESS
    @objc let PORT:UInt16 = Bootstrapper.UDP_PORT
    @objc var socket: GCDAsyncUdpSocket!
    private var bootstrapper:Bootstrapper!
    
    init(bootstrapper: Bootstrapper){
        super.init()
        setupConnection()
        self.bootstrapper = bootstrapper
    }
    
    @objc func setupConnection(){
        socket = GCDAsyncUdpSocket(delegate: self, delegateQueue: DispatchQueue.main)
        socket.setIPv4Enabled(true)
        
        do {
            try socket.bind(toPort: PORT)
            try socket.enableBroadcast(true)
            try socket.joinMulticastGroup(IP)
            try socket.beginReceiving()

        } catch {
            print("Error \(error)")
        }
    }

    public func udpSocket(_ sock: GCDAsyncUdpSocket,
                          didReceive data: Foundation.Data,
                          fromAddress address: Foundation.Data,
                          withFilterContext filterContext: Any?) {
        
        //print("incoming message: \(String(describing: String(data: 
        //      data, encoding: String.Encoding.utf8)))")
        
        var host: NSString?
        var port: UInt16 = 0
        
        if bootstrapper.status != 0x02 {
            switch data[0] {
            case Bootstrapper.DISCOVER_REQUEST:
                GCDAsyncUdpSocket.getHost(&host, port: &port,
                                          fromAddress: address)
                bootstrapper.UDP_CLIENT_ADDRESS = host! as String
                bootstrapper.UDP_CLIENT_PORT = port
                let discoverRequest = bootstrapper.handleDiscoverRequest(client: sock)
                let sendData = Data(bytes: discoverRequest)
                self.socket.send(sendData, toHost: bootstrapper.UDP_CLIENT_ADDRESS,
                                 port: bootstrapper.UDP_CLIENT_PORT,
                                 withTimeout: -1, tag: 0)
                break
                
            case Bootstrapper.PROVISION_REQUEST:
                if bootstrapper.START_PROVISION == false {
                    bootstrapper.START_PROVISION = true
                    let reqArray = [UInt8](data)
                    GCDAsyncUdpSocket.getHost(&host, port: &port,
                                              fromAddress: address)
                    bootstrapper.UDP_CLIENT_ADDRESS = host! as String
                    bootstrapper.UDP_CLIENT_PORT = port
                    bootstrapper.handleProvisionRequest(request: reqArray,
                                                        offset: 1,
                                                        length: data.count)
                }
                break
                
            default:
                break
            }
        } else {
            socket.close()
        }
    }
}
