#!/usr/bin/env python3 import json from base64 import b64encode, b64decode import pyDH from Crypto.Hash import SHA512 from Crypto.Cipher import PKCS1_OAEP from Crypto.Cipher import ChaCha20_Poly1305 from Crypto.PublicKey.RSA import RsaKey from Crypto.Signature import pkcs1_15 from netsim import network_interface from authentication import Authetication class NetWrapper: def __init__(self, clientPublicKey: dict, serverPrivateKey: RsaKey, authenticationInstance: Authetication): self.clientPublicKey = clientPublicKey self.currentClientPublicKey = "".encode('UTF-8') self.serverPrivateKey = serverPrivateKey self.cipherkey = "".encode('UTF-8') self.network = network_interface('./../../netsim/network/', 'A') self.clientAddr = "" self.currentUser = "" self.homeDirectory = "" self.authenticationInstance = authenticationInstance def ecryptRSAMessage(self, message: str) -> bytes: cipher_rsa = PKCS1_OAEP.new(self.currentClientPublicKey) encrypted_msg = cipher_rsa.encrypt(message.encode('UTF-8')) 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.serverPrivateKey).sign(h) return header, headersignature def verifyRSAHeaderSignature(self, header: bytes, headersignature: bytes) -> bool: h = SHA512.new(header) try: pkcs1_15.new(self.currentClientPublicKey).verify(h, headersignature) return True except Exception: return False def decryptRSAMessage(self, message: bytes) -> bytes: cipher_rsa = PKCS1_OAEP.new(self.serverPrivateKey) return cipher_rsa.decrypt(message) def serverIdentify(self, msg: bytes) -> None: incommingJson = json.loads(msg.decode('UTF-8')) header = json.loads(b64decode(incommingJson['header']).decode('UTF-8')) try: if not self.verifyRSAHeaderSignature(b64decode(incommingJson['header']), b64decode(incommingJson['headersignature'])) or header[ 'type'] != 'IDY': raise Exception('Bad initial message') except Exception: raise Exception('Bad initial message') self.clientAddr = header['source'] self.currentUser = header['username'] self.currentClientPublicKey = self.clientPublicKey[self.currentUser] retheader, retheadersignature = self.signRSAHeader("IDY", {}) retmsg = self.ecryptRSAMessage(b64decode(incommingJson['message']).decode('UTF-8')) identMsg = json.dumps( {'header': retheader, 'headersignature': retheadersignature, 'message': b64encode(retmsg).decode('UTF-8')}).encode( 'UTF-8') self.network.send_msg(self.clientAddr, identMsg) def sendMessage(self, message: bytes) -> None: self.sendTypedMessage(message, "CMD") 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({'source': self.network.own_addr, 'type': type}).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') b64tag = b64encode(tag).decode('UTF-8') sendjson = json.dumps( {'header': b64encode(header).decode('UTF-8'), 'nonce': nonce, 'message': ct, 'tag': b64tag}).encode( 'UTF-8') self.network.send_msg(self.clientAddr, sendjson) def keyExchange(self) -> None: dh = pyDH.DiffieHellman() mypubkey = self.ecryptRSAMessage(str(dh.gen_public_key())) 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.clientAddr, jsonmsg) status, msg = self.network.receive_msg(blocking=True) if not status: raise Exception('Network error during connection.') decodedmsg = json.loads(msg.decode('UTF-8')) header = json.loads(b64decode(decodedmsg['header']).decode('UTF-8')) if not self.verifyRSAHeaderSignature(b64decode(decodedmsg['header']), b64decode(decodedmsg['headersignature'])) or not ( header['source'] == self.clientAddr and header['type'] == 'DH'): raise Exception('Header signature error') clientpubkey = int(self.decryptRSAMessage(b64decode(decodedmsg['message'])).decode('UTF-8')) cipherkey = dh.gen_shared_key(clientpubkey).encode('UTF-8') hasher = SHA512.new() hasher.update(cipherkey) self.cipherkey = (hasher.hexdigest()[:32]).encode('UTF-8') def login(self) -> bool: try: status, msg = self.network.receive_msg(blocking=True) if not status: raise Exception('Network error during connection.') cleartext = self.recieveEncryptedMessage(msg, "AUT").decode('UTF-8') if cleartext == "ERROR": return False else: plaintext = cleartext.split(' ') self.homeDirectory = self.authenticationInstance.login(plaintext[1], plaintext[2]) linsuccess = (not (len(plaintext) != 3 or plaintext[0] != "LIN" or plaintext[ 1] != self.currentUser)) and self.homeDirectory if linsuccess: message = "OK".encode('UTF-8') else: message = "ERROR".encode('UTF-8') self.sendTypedMessage(message, "AUT") return linsuccess except Exception: print("Login failed") return False def initClientConnection(self, msg: bytes) -> bytes: print('A client is trying to connect') try: print('Server and Client identity verification started') self.serverIdentify(msg) print('Key exchange started') self.keyExchange() print('Authorization started') success = self.login() print(f'Authorization completed, success: {success}') if success: return "LINOK".encode('UTF-8') else: self.logout() return "LINERROR".encode('UTF-8') except Exception: print("Error ecountered, resetting") self.logout() return "LINERROR".encode('UTF-8') def recieveMessage(self) -> bytes: status, msg = self.network.receive_msg(blocking=True) if not status: raise Exception('Network error during connection.') if not self.clientAddr: return self.initClientConnection(msg) else: return self.recieveEncryptedMessage(msg, "CMD") def logout(self) -> None: self.clientAddr = "" self.cipherkey = "".encode('UTF-8') self.currentClientPublicKey = "".encode('UTF-8') self.currentUser = "" self.homeDirectory = "" def recieveEncryptedMessage(self, msg: bytes, type: str) -> bytes: if not (type == "AUT" or type == "CMD"): raise Exception('Unknown message type') try: b64 = json.loads(msg.decode('UTF-8')) retheader = json.loads(b64decode(b64['header']).decode('UTF-8')) retnonce = b64decode(b64['nonce']) retciphertext = b64decode(b64['message']) rettag = b64decode(b64['tag']) retcipher = ChaCha20_Poly1305.new(key=self.cipherkey, nonce=retnonce) retcipher.update(b64decode(b64['header'])) plaintext = retcipher.decrypt_and_verify(retciphertext, rettag) if not (retheader['source'] == self.clientAddr and retheader['type'] == type): return "ERROR".encode('UTF-8') else: return plaintext except Exception: print("Incorrect decryption") return "ERROR".encode('UTF-8')