#!/usr/bin/env python3
# coding: utf8

import json
import logging
import logging.config
import os
import socket
import threading

try:
    # noinspection PyUnresolvedReferences
    from Crypto.PublicKey import RSA as RSA
    # noinspection PyUnresolvedReferences
    from Crypto.Cipher import PKCS1_OAEP as PKCS1_OAEP

    pycryptodome = False
except ModuleNotFoundError:  # Pycryptodomex
    from Cryptodome.PublicKey import RSA as RSA
    from Cryptodome.Cipher import PKCS1_OAEP as PKCS1_OAEP
    from Cryptodome.Cipher import AES as AES

    pycryptodome = True


class RsaGenThread(threading.Thread):

    def __init__(self, client_, difficulty=2):
        threading.Thread.__init__(self)
        self.client = client_
        self.difficulty = difficulty

    def run(self):
        if os.path.isfile("private.pem"):
            try:
                with open("private.pem", "rb") as keyfile:
                    rsa = RSA.importKey(keyfile.read())
                    if not rsa.has_private():
                        warning("Le fichier clef ne contient pas de clef privée")
                        raise ValueError
                    self.client.rsa = rsa
            except (IndexError, ValueError):
                warning("Fichier clef corrompu")
                debug("Suppression du fichier clef corromu")
                os.remove("private.pem")
        if not os.path.isfile("private.pem"):  # We're not using if/else because we may delete the file in the
            # previous if statement
            debug("Génération de la clef RSA pour %s" % self.client.name)
            self.client.rsa = RSA.generate(BUFFER_SIZE + 256 * self.difficulty)
            with open("private.pem", "wb") as keyfile:
                keyfile.write(self.client.rsa.exportKey())
            with open("public.pem", "wb") as keyfile:
                keyfile.write(self.client.rsa.publickey().exportKey())


def setup_logging(default_path='log_config.json', default_level=logging.INFO, env_key='LOG_CFG'):
    """Setup logging configuration
    """
    path = default_path
    value = os.getenv(env_key, None)
    if value:
        path = value
    if os.path.exists(path):
        with open(path, 'rt') as f:
            config = json.load(f)
        logging.config.dictConfig(config)
    else:
        logging.basicConfig(level=default_level)


setup_logging()

log_server = logging.getLogger('server')

debug = log_server.debug
info = log_server.info
warning = log_server.warning
error = log_server.error
critical = log_server.critical

#### Variables ####
HOST = '127.0.0.1'
PORT = 8888
BUFFER_SIZE = 4096
CHUNK_SIZE = int(BUFFER_SIZE / 8)
BEGIN_MESSAGE = bytes("debut".ljust(BUFFER_SIZE, ";"), "ascii")
END_MESSAGE = bytes("fin".ljust(BUFFER_SIZE, ";"), "ascii")
NAME = b"client1"
VERSION = b"EICP2P2 V1"

REQUEST_TYPE = [
    b'ping', b'pingACK', b'updateAsk', b'updateBack', b'transfer', b'register_client', b'registerACK', b'send', b'sendACK',
    b'exit', b'RSASend', b'init', b'getUsers', b'getUsersACK',
]

clientSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
clientSocket.connect((HOST, PORT))


# ET ICI ON MET LE CLIENT

class MainThread(threading.Thread):
    """Main client class, for each identity."""

    def __init__(self, name):
        threading.Thread.__init__(self)
        self.name = name
        self.rsa = None
        self.aes_key = None
        self.clientSocket = None
        self.users = {}

    def run(self):
        rsa = RsaGenThread(self)
        rsa.start()
        rsa.join()
        self.clientSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.clientSocket.connect((HOST, PORT))
        self.initialisation()
        self.register()
        self.get_users()
        print(self.users)
        while True:
            try:
                type = bytes(input("Type: "), "ascii")
                contenu = bytes(input("Contenu: "), "ascii")
                self.send_aes(self.gen_header(type_=type)+contenu)
                print(self.receive_aes())
            except Exception as e:
                print(e)

    def register(self):
        self.send_aes(self.gen_header(type_=b"register_client"))
        self.id_noeud, self.id_client = self.receive_aes()[BUFFER_SIZE:].split(b"{%=&%&=%}")

    def initialisation(self):
        header = self.gen_header(b"RSASend", from_=b"client")
        content = self.rsa.publickey().exportKey()
        self.send(header + content)
        print("oki1")
        data = self.receive_rsa()
        print("oki2")
        header = self.extract_header(data[:BUFFER_SIZE])
        if header[b"type"] != b"init":
            debug("Incorrect request type, end of connection")
            return
        self.aes_key = data[BUFFER_SIZE:]
        debug("End of initialisation")

    @staticmethod
    def extract_header(data):
        """Extract header from data

        :param data: Data to extract header
        :type data: bytes

        :return: Dictionary with header datas
        :rtype: dict{bytes: bytes}"""
        if len(data) > BUFFER_SIZE:
            debug("Header too long")
            data = data[:BUFFER_SIZE]
        data_lines = data.split(b'\n')
        if data_lines[0] != VERSION:
            raise ValueError("Version is incorrect.")
        return {
            l.split(b": ")[0]: l.split(b": ")[1].rstrip(b";") for l in data_lines[1:]
        }

    @staticmethod
    def gen_header(type_, to_=None, from_=None):
        """Generate header

        :param type_: Request type
        :param to_: `to` field in header, cf ../RFC8497.md
        :param from_: `from` field in header, cf ../RFC8497.md
        :type type_: bytes
        :type to_: bytes
        :type from_: bytes

        :return: header
        :rtype: bytes"""
        if type_ not in REQUEST_TYPE:
            raise ValueError("Unknown request type")
        header = VERSION + b"\ntype: " + type_
        if to_:
            header += b"\nto: " + to_
        if from_:
            header += b"\nfrom: " + from_
        return header.ljust(BUFFER_SIZE, b';')

    def get_users(self):
        self.send_aes(self.gen_header(type_=b'getUsers'))
        back = self.receive_aes()
        print(back)
        header = self.extract_header(back)
        contenu = back[BUFFER_SIZE:]
        if header[b"type"] == b"getUsersACK":
            self.users = [(contenu.split(b"%!!%")[i], i) for i in range(len(contenu.split(b"%!!%")))]

    def send_aes(self, to_send, key=None):
        """Send message with aes encryption

        :param to_send: Message to send
        :type to_send: bytes
        :param key: key to replace self.aes_key
        :type key: bytes

        :rtype: NoneType
        :return: Nothing"""
        debug(b"Send with AES encryption: " + to_send)
        if key is None:
            key = self.aes_key
        if key is None:
            info("AES key not generated, connection failure.")
            self.client.send(b"Error")
            return
        # Get RSA key
        aes_object = AES.new(key, AES.MODE_ECB)
        encrypted = b""
        for to_send_text in [to_send[i:i + 32] for i in range(0, len(to_send), 32)]:
            encrypted += aes_object.encrypt(to_send_text.ljust(32, b"\x00"))
        self.send(encrypted)
        return None

    def receive_aes(self, key=None):
        """Receive message with aes encryption

        :param key: key to replace self.aes_key
        :type key: bytes
        """
        to_decrypt = self.receive()
        debug(b"Received in aes "+to_decrypt)
        if key is None:
            key = self.aes_key
        if key is None:
            info("AES key not generated, connection failure.")
            self.client.send(b"Error")
            return
        aes_object = AES.new(key, AES.MODE_ECB)
        decrypted = b""
        for block in [to_decrypt[i:i + 32] for i in range(0, len(to_decrypt), 32)]:
            print(len(block))
            decrypted += aes_object.decrypt(block)
        return decrypted.rstrip(b"\x00")

    ################################################ COMMUNICATION WITH RSA ############################################

    def receive_rsa(self, rsa_key=None):
        if rsa_key is None:
            rsa_key = self.rsa.exportKey()
        if rsa_key is None:
            info("RSA key not generated, connection failure.")
            return
        recipient_key = RSA.importKey(rsa_key)
        cipher_rsa = PKCS1_OAEP.new(recipient_key)
        raw = self.receive()
        content = b""
        for to_decrypt in [raw[i:i + self.rsa.publickey().size_in_bytes()] for i in
                           range(0, len(raw), self.rsa.publickey().size_in_bytes())]:
            content += cipher_rsa.decrypt(to_decrypt)
        return content

    ############################################ COMMUNICATION WITHOUT CRYPTING ############################################

    def receive(self):
        """Receive message from connection

        :rtype: bytes
        :return: Message's content"""
        chunk = bytes("", "ascii")  # Temp variable to store received datas
        while chunk != BEGIN_MESSAGE:
            chunk = self.clientSocket.recv(BUFFER_SIZE)
        content = b''
        while chunk != END_MESSAGE:
            chunk = self.clientSocket.recv(BUFFER_SIZE)
            # Get only interesting chucks
            if chunk != END_MESSAGE:
                # Get content part
                # int.from_bytes(chunk[:2], byteorder='big') == Get content size
                content += chunk[2:int.from_bytes(chunk[:2], byteorder='big') + 2]
        debug(b"Received: " + content)
        return content

    def send(self, to_send):
        """Send message to connection

        :param to_send: message to send
        :type to_send: bytes

        :return: Nothing
        :rtype: NoneType"""
        debug(b"Send " + to_send)
        # Sending the message start
        self.clientSocket.send(BEGIN_MESSAGE)
        i = 0
        for to_send_text in [to_send[i:i + BUFFER_SIZE - 2] for i in range(0, len(to_send), BUFFER_SIZE - 2)]:
            self.clientSocket.send(
                (len(to_send_text)).to_bytes(2, byteorder='big')  # Size of the message contained by the chunk
                + to_send_text.ljust(BUFFER_SIZE - 2, bytes(1))  # Content of the chunk
            )
            i += 1
        # Sending the message stop
        self.clientSocket.send(END_MESSAGE)
        return None


if __name__ == "__main__":
    client = MainThread(NAME)
    client.start()