#!/usr/bin/env python3
import json
from base64 import b64encode, b64decode
import pyDH
from Crypto.Hash import SHA512
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Cipher import ChaCha20_Poly1305
from Crypto.PublicKey.RSA import RsaKey

from netsim import network_interface
from authentication import Authetication


class NetWrapper:

    def __init__(self, clientPublicKey: dict, serverPrivateKey: RsaKey, authenticationInstance: Authetication):
        self.clientPublicKey = clientPublicKey
        self.currentClientPublicKey = "".encode('UTF-8')
        self.serverPrivateKey = serverPrivateKey
        self.cipherkey = "".encode('UTF-8')
        self.network = network_interface('./../../netsim/network/', 'A')
        self.clientAddr = ""
        self.currentUser = ""
        self.homeDirectory = ""
        self.authenticationInstance = authenticationInstance

    def serverIdentify(self, msg: bytes) -> None:
        incommingJson = json.loads(msg.decode('UTF-8'))
        if incommingJson['type'] != "IDY":
            raise Exception('Wrong message type encountered')
        self.clientAddr = incommingJson['source']
        self.currentUser = incommingJson['username']
        self.currentClientPublicKey = self.clientPublicKey[self.currentUser]
        cipher_rsa = PKCS1_OAEP.new(self.serverPrivateKey)
        retmsg = cipher_rsa.decrypt(b64decode(incommingJson['message'])).decode('UTF-8')
        cipher = PKCS1_OAEP.new(self.currentClientPublicKey)

        identMsg = json.dumps(
            {'type': 'IDY', 'source': self.network.own_addr,
             'message': b64encode(cipher.encrypt(retmsg.encode('UTF-8'))).decode('UTF-8')}).encode(
            'UTF-8')
        self.network.send_msg(self.clientAddr, identMsg)

    def sendMessage(self, message: bytes) -> None:
        self.sendTypedMessage(message, "CMD")

    def sendTypedMessage(self, message: bytes, type: str) -> None:
        if not (type == "IDY" or type == "DH" or type == "CMD"):
            raise Exception('Unknown message type')
        cipher = ChaCha20_Poly1305.new(key=self.cipherkey)
        header = json.dumps({'source': self.network.own_addr, 'type': type}).encode('UTF-8')
        cipher.update(header)
        ciphertext, tag = cipher.encrypt_and_digest(message)
        nonce = b64encode(cipher.nonce).decode('UTF-8')
        ct = b64encode(ciphertext).decode('UTF-8')
        b64tag = b64encode(tag).decode('UTF-8')
        sendjson = json.dumps(
            {'header': b64encode(header).decode('UTF-8'), 'nonce': nonce, 'message': ct, 'tag': b64tag}).encode(
            'UTF-8')
        self.network.send_msg(self.clientAddr, sendjson)

    def keyExchange(self) -> None:
        dh = pyDH.DiffieHellman()
        cipher = PKCS1_OAEP.new(self.currentClientPublicKey)
        mypubkey = b64encode(cipher.encrypt(str(dh.gen_public_key()).encode('UTF-8'))).decode('UTF-8')
        jsonmsg = json.dumps({'type': 'DH', 'source': self.network.own_addr, 'message': mypubkey}).encode('UTF-8')
        self.network.send_msg(self.clientAddr, jsonmsg)
        decodedmsg = {'source': '', 'type': ''}
        while not (decodedmsg['source'] == self.clientAddr and decodedmsg['type'] == 'DH'):
            status, msg = self.network.receive_msg(blocking=True)
            if not status:
                raise Exception('Network error during connection.')
            decodedmsg = json.loads(msg.decode('UTF-8'))
        cipher_rsa = PKCS1_OAEP.new(self.serverPrivateKey)
        serverpubkey = int(cipher_rsa.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
        cipherkey = dh.gen_shared_key(serverpubkey).encode('UTF-8')
        hasher = SHA512.new()
        hasher.update(cipherkey)
        self.cipherkey = (hasher.hexdigest()[:32]).encode('UTF-8')

    def login(self) -> bool:
        try:
            status, msg = self.network.receive_msg(blocking=True)
            if not status:
                raise Exception('Network error during connection.')
            cleartext = self.recieveEncryptedMessage(msg, "AUT").decode('UTF-8')
            if cleartext=="ERROR":
                return False
            else:
                plaintext = cleartext.split(' ')
            self.homeDirectory = self.authenticationInstance.login(plaintext[1], plaintext[2])
            linsuccess = (not (len(plaintext) != 3 or plaintext[0] != "LIN" or plaintext[
                1] != self.currentUser)) and self.homeDirectory
            if linsuccess:
                message = "OK".encode('UTF-8')
            else:
                message = "ERROR".encode('UTF-8')
            self.sendTypedMessage(message, "AUT")
            return linsuccess
        except Exception:
            print("Login failed")
            return False

    def initClientConnection(self, msg: bytes) -> bytes:
        print('3')
        try:
            print('4')
            self.serverIdentify(msg)
            print('5')
            self.keyExchange()
            print('6')
            success = self.login()
            print('7')
            if success:
                return "LINOK".encode('UTF-8')
            else:
                self.logout()
                return "LINERROR".encode('UTF-8')
        except Exception:
            self.logout()
            return "LINERROR".encode('UTF-8')

    def recieveMessage(self) -> bytes:
        status, msg = self.network.receive_msg(blocking=True)
        if not status:
            raise Exception('Network error during connection.')
        if not self.clientAddr:
            return self.initClientConnection(msg)
        else:
            return self.recieveEncryptedMessage(msg, "CMD")

    def logout(self) -> None:
        self.clientAddr = ""
        self.cipherkey = "".encode('UTF-8')
        self.currentClientPublicKey = "".encode('UTF-8')
        self.currentUser = ""
        self.homeDirectory = ""

    def recieveEncryptedMessage(self, msg: bytes, type: str) -> bytes:
        if not (type == "IDY" or type == "DH" or type == "CMD"):
            raise Exception('Unknown message type')
        try:
            b64 = json.loads(msg.decode('UTF-8'))
            retheader = json.loads(b64decode(b64['header']).decode('UTF-8'))
            retnonce = b64decode(b64['nonce'])
            retciphertext = b64decode(b64['message'])
            rettag = b64decode(b64['tag'])
            retcipher = ChaCha20_Poly1305.new(key=self.cipherkey, nonce=retnonce)
            retcipher.update(b64decode(b64['header']))
            plaintext = retcipher.decrypt_and_verify(retciphertext, rettag)
            if not (retheader['source'] == self.clientAddr and retheader['type'] == type):
                return "ERROR".encode('UTF-8')
            else:
                return plaintext
        except Exception:
            print("Incorrect decryption")
            return "ERROR".encode('UTF-8')