"""
A middleman IoT websocket device which decrypts and re-encrypts a message

@author hornetfighter515
@author Lachezar Todorov
"""
import websocket
import _thread
import time
import yaml
import json 
import rsa


def send(ws, message):
    try:
        ws.send(json.dumps(message))
    except:
        print(f'Sending of message {message} failed.')


def on_message(ws, message):
    # decrypt the message
    #try:
    if True:
        msg = json.loads(message)
        contents = msg['contents'].encode('latin')
        res = {
            "from":msg['to'],
            "to":0,
            "type":'message',
            "id":msg['id'],
            "contents": rsa.decrypt(contents, priv).decode('latin')
            }
        print('Forwarding message...')
        send(ws, res)
    #except:
    #    print('Sending failed.')
    #    send(ws, message)


def on_error(ws, error):
    print(error)

def on_close(ws, close_status_code, close_msg):
    print(f"### closed. reason: {close_msg} ###")

def on_open(ws):
    def run(*args):
        send(ws, {
            "pub_key":{
                "n":pub.n,
                "e":pub.e
                }
            })
        running = True
        while running:
            time.sleep(1)
    _thread.start_new_thread(run, ())

def open_socket(url, port):
    ws = websocket.WebSocketApp(f"ws://{url}:{port}",
        on_open=on_open,
        on_message=on_message,
        on_error=on_error,
        on_close=on_close)
    ws.run_forever()


if __name__ == "__main__":
    with open('src/config.yml', 'r') as f:
        conf = yaml.safe_load(f)
    global pub, priv
    # eventually put this in config
    print('Generating keys...')
    pub, priv = rsa.newkeys(2048)
    print('Keys generated!')
    open_socket(conf['ws']['iot']['url'], conf['ws']['iot']['port'])