161 lines
7.3 KiB
Python
161 lines
7.3 KiB
Python
#!/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, serverPubKey: RsaKey, username: str = "",
|
|
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()
|
|
encrypted_msg = self.ecryptRSAMessage(randommsg.encode('UTF-8'))
|
|
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': b64encode(mypubkey).decode('UTF-8')}).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:
|
|
if self.username == "":
|
|
raise Exception('Username is not initialized')
|
|
if password == "":
|
|
raise Exception('Passowrd may not be empty')
|
|
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')
|