client/client/netwrapper.py

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

from netsim import network_interface


class NetWrapper:
    def __init__(self, privateKey: RsaKey, clientAddress: str, username: str, serverPubKey: RsaKey,
                 serverAddr: str = 'A'):
        self.network = network_interface('./../../netsim/network/', clientAddress)
        self.serverAddr = serverAddr
        self.username = username
        self.privateKey = privateKey
        self.serverPubKey = serverPubKey
        self.cipherkey = "".encode('UTF-8')

    def randomStringGenerator(self, str_size: int = 256,
                              allowed_chars: str = string.ascii_letters + string.punctuation) -> str:
        return ''.join(random.choice(allowed_chars) for x in range(str_size))

    def ecryptRSAMessage(self, message: bytes) -> bytes:
        cipher_rsa = PKCS1_OAEP.new(self.serverPubKey)
        encrypted_msg = cipher_rsa.encrypt(message)
        return encrypted_msg

    def signRSAHeader(self, type: str, extradata: dict) -> (bytes, bytes):
        mandatory = {'type': type, 'source': self.network.own_addr}
        header = json.dumps({**mandatory, **extradata}).encode('UTF-8')
        h = SHA512.new(header)
        headersignature = pkcs1_15.new(self.privateKey).sign(h)
        return header, headersignature

    def verifyRSAHeaderSignature(self, header: bytes, headersignature: bytes) -> bool:
        h = SHA512.new(header)
        try:
            pkcs1_15.new(self.serverPubKey).verify(h, headersignature)
            return True
        except Exception:
            return False

    def decryptRSAMessage(self, message: bytes) -> bytes:
        cipher_rsa = PKCS1_OAEP.new(self.privateKey)
        return cipher_rsa.decrypt(message)

    def recieveAndUnpackRSAMessage(self) -> (str, str):
        status, msg = self.network.receive_msg(blocking=True)
        if not status:
            raise Exception('Network error during connection.')
        decodedmsg = json.loads(msg.decode('UTF-8'))
        return decodedmsg, json.loads(b64decode(decodedmsg['header']).decode('UTF-8'))

    def identifyServer(self) -> bool:
        randommsg = self.randomStringGenerator().encode('UTF-8')
        encrypted_msg = self.ecryptRSAMessage(randommsg)
        header, headersignature = self.signRSAHeader('IDY', {'username': self.username})
        identMsg = json.dumps(
            {'header': b64encode(header).decode('UTF-8'),
             'message': b64encode(encrypted_msg).decode('UTF-8'),
             'headersignature': b64encode(headersignature).decode('UTF-8')}).encode(
            'UTF-8')
        self.network.send_msg(self.serverAddr, identMsg)
        returnJson, header = self.recieveAndUnpackRSAMessage()
        try:
            if not self.verifyRSAHeaderSignature(b64decode(returnJson['header']),
                                                 b64decode(returnJson['headersignature'])) or not (
                    header['source'] == self.serverAddr and header['type'] == 'IDY'):
                return False
            retmsg = self.decryptRSAMessage(b64decode(returnJson['message'])).decode('UTF-8')
        except Exception:
            return False
        return retmsg == randommsg

    def createEncryptedChannel(self) -> None:
        dh = pyDH.DiffieHellman()
        mypubkey = self.ecryptRSAMessage(str(dh.gen_public_key()).encode('UTF-8'))
        header, headersignature = self.signRSAHeader("DH",{})
        jsonmsg = json.dumps(
            {'header': b64encode(header).decode('UTF-8'), 'headersignature': b64encode(headersignature).decode('UTF-8'),
             'message': mypubkey}).encode('UTF-8')
        self.network.send_msg(self.serverAddr, jsonmsg)
        decodedmsg, header = self.recieveAndUnpackRSAMessage()
        if not self.verifyRSAHeaderSignature(b64decode(decodedmsg['header']),
                                             b64decode(decodedmsg['headersignature'])) or not (
                header['source'] == self.serverAddr and header['type'] == 'DH'):
            raise Exception('Header signature error')
        serverpubkey = int(self.decryptRSAMessage(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 authenticate(self, password: str) -> None:
        message = f"LIN {self.username} {password}".encode('UTF-8')
        self.sendTypedMessage(message, "AUT")
        plaintext = self.recieveTypedMessage("AUT").decode('UTF-8')
        if plaintext != "OK":
            print("Authentication error")

    def connectToServer(self, password: str) -> None:
        identStatus = self.identifyServer()
        if not identStatus:
            raise Exception('Server identification faliure')
        self.createEncryptedChannel()
        self.authenticate(password)

    def sendMessage(self, message: bytes) -> None:
        try:
            self.sendTypedMessage(message, "CMD")
        except Exception:
            print("Could not send message")

    def sendTypedMessage(self, message: bytes, type: str) -> None:
        if not (type == "AUT" or type == "CMD"):
            raise Exception('Unknown message type')
        cipher = ChaCha20_Poly1305.new(key=self.cipherkey)
        header = json.dumps({'type': type, 'source': self.network.own_addr}).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')
        sendjson = json.dumps({'header': b64encode(header).decode('UTF-8'), 'nonce': nonce, 'message': ct,
                               'tag': b64encode(tag).decode('UTF-8')}).encode(
            'UTF-8')
        self.network.send_msg(self.serverAddr, sendjson)

    def recieveMessage(self) -> bytes:
        return self.recieveTypedMessage("CMD")

    def recieveTypedMessage(self, type: str) -> bytes:
        if not (type == "AUT" or type == "CMD"):
            raise Exception('Unknown message type')
        try:
            status, msg = self.network.receive_msg(blocking=True)
            if not status:
                raise Exception('Network error during connection.')
            b64 = json.loads(msg.decode('UTF-8'))
            retnonce = b64decode(b64['nonce'])
            retciphertext = b64decode(b64['message'])
            retcipher = ChaCha20_Poly1305.new(key=self.cipherkey, nonce=retnonce)
            retcipher.update(b64decode(b64['header']))
            plaintext = retcipher.decrypt_and_verify(retciphertext, b64decode(b64['tag']))
            retheader = json.loads(b64decode(b64['header']).decode('UTF-8'))
            if not (retheader['source'] == self.serverAddr and retheader['type'] == type):
                return "ERROR".encode('UTF-8')
            return plaintext
        except Exception:
            print("Incorrect decryption")
            return "ERROR".encode('UTF-8')