secure-chat/server/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
    # noinspection PyUnresolvedReferences,PyProtectedMember
    from Crypto.Random._UserFriendlyRNG import get_random_bytes as get_random_bytes

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

    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 = ''
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")

VERSION = "EICP2P2 V1"

DONE = 0
ERROR = 1

#### Socket ####
main_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
main_socket.bind((HOST, PORT))


#### Threads ####
class ClientThread(threading.Thread):
    """Main thread, for each client"""

    def __init__(self, clientsocket, ip_client, port):
        """Create ClientThread object

        :param clientsocket: Client's socket
        :param ip_client: Client's ip address
        :param port: Client's connection PORT
        :type clientsocket: socket.socket
        :type ip_client: str
        :type port: int

        :return: Nothing
        :rtype: NoneType"""
        debug("Creation du thread pour %s" % ip_client)
        threading.Thread.__init__(self)  # initialisation du thread
        self.client = clientsocket
        self.ip = ip_client
        self.port = port
        self.running = True
        self.status = None
        self.rsa_client = None
        self.aes_key = get_random_bytes(64)
        debug("Creation du thread pour %s reussie" % ip_client)

    def initialize(self):
        """Initialize connection with client

        :rtype: NoneType
        :return: Nothing"""
        # Receive message
        message = self.receive()
        # Get informations
        header = message[:BUFFER_SIZE]
        content = message[BUFFER_SIZE:]
        # Extract header data
        header_txt = header.decode("ascii").rstrip(";")
        version = header_txt.split("\n")[0]
        # Extraction des champs du header
        if version != VERSION:
            debug("Incorrect protocol version.")
            self.send(b"Error")
        # Extract information from header
        champs = {l.split(": ")[0]: l.split(": ")[1] for l in header_txt.split('\n')[1:]}

        if champs.get("type", None) is None:
            debug("The type field is not in the header")
            self.send(b"Error")
            return

        if self.status is None and champs["type"] != "RSASend":
            debug("Requête différente de RASSend avec une connection non initialisée")
            self.send(b"Error".ljust(BUFFER_SIZE, b";"))
            return

        if champs["type"] == "RSASend":
            debug("Réception de la clef RSA de %s", self.ip)
            self.rsa_client = content.rstrip(b";").rstrip(b' ')
            header = bytes((VERSION + "\ntype: init").ljust(BUFFER_SIZE, ";"), 'ascii')
            content = self.aes_key.ljust(BUFFER_SIZE, b";")
            self.send_rsa(header + content)
            return

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

    def send_rsa(self, to_send, key=None):
        """Send message with rsa encryption

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

        :rtype: NoneType
        :return: Nothing"""
        debug(b"Send with RSA encryption: " + to_send + bytes(str(self.ip), "ascii"))
        if key is None:
            key = self.rsa_client
        if key is None:
            info("RSA key not received, connection failure.")
            self.client.send(b"Error".ljust(BUFFER_SIZE, bytes(1)))
            return
        # Get RSA key
        print(len(to_send))
        recipient_key = RSA.importKey(key)
        # RSA encryption object
        cipher_rsa = PKCS1_OAEP.new(recipient_key)
        encrypted = b""
        for to_send_text in [to_send[i:i + CHUNK_SIZE] for i in range(0, len(to_send), CHUNK_SIZE)]:
            print(len(to_send_text))
            encrypted += cipher_rsa.encrypt(to_send_text)
            print(len(encrypted))
        self.send(encrypted)
        return None

    ############################################ 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.client.recv(BUFFER_SIZE)
        content = b''
        while chunk != END_MESSAGE:
            chunk = self.client.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')]
        debug(b"Received from" + bytes(str(self.ip), ascii) + b" : " + 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 + b" to " + bytes(str(self.ip), "ascii"))
        # Sending the message start
        self.client.send(BEGIN_MESSAGE)
        i = 0
        print(len(to_send))
        for to_send_text in [to_send[i:i + BUFFER_SIZE - 2] for i in range(0, len(to_send), BUFFER_SIZE - 2)]:
            self.client.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.client.send(END_MESSAGE)
        return None

    def run(self):  # main de la connection du client
        """Run thread mainloop

        :return: Nothing
        :rtype: NoneType"""
        info(self.ip + "connected, initialize connection...")
        self.initialize()
        info(self.ip + "connection initialized.")
        print("done")
        self.client.close()


if __name__ == "__main__":
    clients = []
    while True:
        main_socket.listen(1)  # Waiting for incoming connections
        client_socket, (ip, PORT) = main_socket.accept()
        newClient = ClientThread(client_socket, ip, PORT)
        newClient.start()
        clients.append(newClient)