secure-chat/client/main.py

#!/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

    pycryptodome = True


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', b'registerACK', b'send', b'sendACK',
    b'exit', b'RSASend', b'init',
]

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

    def run(self):
        rsa = self.RsaGenThread(self)
        rsa.start()
        rsa.join()
        self.clientSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.clientSocket.connect((HOST, PORT))
        self.initialisation()

    def initialisation(self):
        header = self.gen_header(b"RSASend", from_=b"client")
        content = self.rsa.publickey().exportKey()
        self.send(header + content)
        data = self.receive_rsa()
        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';')

    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())

    ################################################ 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()