refactor chacha
All checks were successful
continuous-integration/drone/push Build is passing

This commit is contained in:
Torma Kristóf 2021-04-26 18:31:50 +02:00
parent ce4fd0cd30
commit 78d597425d
Signed by: tormakris
GPG Key ID: DC83C4F2C41B1047

View File

@ -7,7 +7,6 @@ import pyDH
from Crypto.Hash import SHA512
from Crypto.Cipher import ChaCha20_Poly1305, PKCS1_OAEP
from Crypto.PublicKey.RSA import RsaKey
from Crypto.Random import get_random_bytes
from netsim import network_interface
@ -66,32 +65,10 @@ class NetWrapper:
def authenticate(self, password: str):
message = f"LIN {self.username} {password}".encode('UTF-8')
cipher = ChaCha20_Poly1305.new(key=self.cipherkey)
header = json.dumps({'type': 'AUT', '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')
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.serverAddr, sendjson)
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']))
retheader = json.loads(b64decode(b64['header']).decode('UTF-8'))
plaintext = retcipher.decrypt_and_verify(retciphertext, b64decode(b64['tag'])).decode('UTF-8')
if plaintext != "OK" or not (retheader['source'] == self.serverAddr and retheader['type'] == 'AUT'):
raise Exception('Authentication error')
except Exception as e:
print(e)
self.sendTypedMessage(message, "AUT")
plaintext = self.recieveTypedMessage("AUT").decode('UTF-8')
if plaintext != "OK":
print("Authentication error")
def connectToServer(self):
identStatus = self.identifyServer()
@ -102,9 +79,17 @@ class NetWrapper:
pw = input()
self.authenticate(pw)
def sendMessage(self, message: bytes):
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': 'CMD', 'source': self.network.own_addr}).encode('UTF-8')
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')
@ -115,6 +100,11 @@ class NetWrapper:
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:
@ -126,7 +116,7 @@ class NetWrapper:
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'] == 'CMD'):
if not (retheader['source'] == self.serverAddr and retheader['type'] == type):
return "ERROR".encode('UTF-8')
return plaintext
except Exception: