187 lines
8.2 KiB
Python
187 lines
8.2 KiB
Python
#!/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 encryptRSAMessage(self, message: bytes) -> bytes:
|
|
cipher_rsa = PKCS1_OAEP.new(self.currentClientPublicKey)
|
|
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.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'))
|
|
self.clientAddr = header['source']
|
|
self.currentUser = header['username']
|
|
self.currentClientPublicKey = self.clientPublicKey[self.currentUser]
|
|
if not self.verifyRSAHeaderSignature(b64decode(incommingJson['header']),
|
|
b64decode(incommingJson['headersignature'])) or header[
|
|
'type'] != 'IDY':
|
|
raise Exception('Bad initial message')
|
|
retheader, retheadersignature = self.signRSAHeader("IDY", {})
|
|
dcryptedmsg = self.decryptRSAMessage(b64decode(incommingJson['message']))
|
|
retmsg = self.encryptRSAMessage(dcryptedmsg)
|
|
identMsg = json.dumps(
|
|
{'header': b64encode(retheader).decode('UTF-8'),
|
|
'headersignature': b64encode(retheadersignature).decode('UTF-8'),
|
|
'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.encryptRSAMessage(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': b64encode(mypubkey).decode('UTF-8')}).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:
|
|
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
|
|
|
|
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')
|