secure-chat/client/main.py

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

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


#### logging ####
# json decoder for int keys
import threading

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP


class Decoder(json.JSONDecoder):
    def decode(self, s, **kwargs):
        result = super().decode(s)  # result = super(Decoder, self).decode(s) for Python 2.x
        return self._decode(result)

    def _decode(self, o):
        if isinstance(o, str):
            try:
                return int(o)
            except ValueError:
                return o
        elif isinstance(o, dict):
            return {k: self._decode(v) for k, v in o.items()}
        elif isinstance(o, list):
            return [self._decode(v) for v in o]
        else:
            return o


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)
NAME = 'default'


# 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

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

    class rsaGenThread(threading.Thread):

        def __init__(self, client, difficulty=30, new=False):
            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())


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