118 lines
5.2 KiB
Python
118 lines
5.2 KiB
Python
#!/usr/bin/env python3
|
|
import random
|
|
import string
|
|
import json
|
|
from base64 import b64encode, b64decode
|
|
import pyDH
|
|
from Crypto.Cipher import PKCS1_OAEP
|
|
from Crypto.PublicKey import RSA
|
|
from Crypto.Cipher import ChaCha20
|
|
from Crypto.Random import get_random_bytes
|
|
|
|
from netsim import network_interface
|
|
|
|
|
|
class NetWrapper:
|
|
SERVER_ADDRESS = 'A'
|
|
|
|
def __init__(self, privateKey: str, clientAddress: str, username: str, serverPubKey: str = "",
|
|
serverAddr: str = 'A'):
|
|
# Create network_interface: network_interface(path, addr) path root is shared with network / addr is own address
|
|
self.network = network_interface('./', clientAddress)
|
|
self.serverAddr = serverAddr
|
|
self.username = username
|
|
self.privateKey = privateKey
|
|
self.serverPubKey = serverPubKey
|
|
self.cipherkey="".encode('UTF-8')
|
|
|
|
def randomStringGenerator(self, str_size: int = 512,
|
|
allowed_chars: str = string.ascii_letters + string.punctuation) -> str:
|
|
return ''.join(random.choice(allowed_chars) for x in range(str_size))
|
|
|
|
def identifyServer(self) -> bool:
|
|
randommsg = self.randomStringGenerator()
|
|
rsakey = RSA.import_key(self.serverPubKey)
|
|
cipher = PKCS1_OAEP.new(rsakey)
|
|
identMsg = json.dumps(
|
|
{'type': 'IDY', 'source': self.network.own_addr, 'username': self.username,
|
|
'message': b64encode(cipher.encrypt(randommsg.encode('UTF-8')))}).encode(
|
|
'UTF-8')
|
|
self.network.send_msg(self.serverAddr, identMsg)
|
|
status, msg = self.network.receive_msg(blocking=True)
|
|
if not status:
|
|
raise Exception('Network error during connection.')
|
|
myrsakey = RSA.import_key(self.privateKey)
|
|
mycipher = PKCS1_OAEP.new(myrsakey)
|
|
returnJson = json.loads(msg.decode('UTF-8'))
|
|
retmsg = mycipher.decrypt(b64decode(returnJson['message'])).decode('UTF-8')
|
|
return retmsg == randommsg
|
|
|
|
def createEncryptedChannel(self):
|
|
dh = pyDH.DiffieHellman()
|
|
rsakey = RSA.import_key(self.serverPubKey)
|
|
cipher = PKCS1_OAEP.new(rsakey)
|
|
mypubkey = b64encode(cipher.encrypt(str(dh.gen_public_key()).encode('UTF-8')))
|
|
jsonmsg = json.dumps({'type': 'DH', 'source': self.network.own_addr, 'message': mypubkey}).encode('UTF-8')
|
|
self.network.send_msg(self.serverAddr,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'))
|
|
myrsakey = RSA.import_key(self.privateKey)
|
|
mycipher = PKCS1_OAEP.new(myrsakey)
|
|
serverpubkey = int(mycipher.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
|
|
self.cipherkey = dh.gen_shared_key(serverpubkey).encode('UTF-8')
|
|
|
|
def authenticate(self, password: str):
|
|
message = f"LIN {self.username} {password}".encode('UTF-8')
|
|
cipher = ChaCha20.new(self.cipherkey, get_random_bytes(12))
|
|
ciphertext = cipher.encrypt(message)
|
|
nonce = b64encode(cipher.nonce).decode('UTF-8')
|
|
ct = b64encode(ciphertext).decode('UTF-8')
|
|
sendjson = json.dumps({'type': 'AUT', 'source': self.network.own_addr, 'nonce': nonce, 'message': ct}).encode('UTF-8')
|
|
self.network.send_msg(self.serverAddr, sendjson)
|
|
status, msg = self.network.receive_msg(blocking=True)
|
|
if not status:
|
|
raise Exception('Network error during connection.')
|
|
try:
|
|
b64 = json.loads(msg)
|
|
retnonce = b64decode(b64['nonce'])
|
|
retciphertext = b64decode(b64['message'])
|
|
retcipher = ChaCha20.new(self.cipherkey, nonce=retnonce)
|
|
plaintext = retcipher.decrypt(retciphertext).decode('UTF-8')
|
|
if plaintext != "OK":
|
|
raise Exception('Authentication error')
|
|
except Exception:
|
|
print("Incorrect decryption")
|
|
|
|
def connectToServer(self):
|
|
identStatus = self.identifyServer()
|
|
if not identStatus:
|
|
raise Exception('Server identification faliure')
|
|
self.createEncryptedChannel()
|
|
print('Please enter your password:')
|
|
pw = input()
|
|
self.authenticate(pw)
|
|
|
|
def sendMessage(self, message: bytes):
|
|
cipher = ChaCha20.new(self.cipherkey, get_random_bytes(12))
|
|
ciphertext = cipher.encrypt(message)
|
|
nonce = b64encode(cipher.nonce).decode('UTF-8')
|
|
ct = b64encode(ciphertext).decode('UTF-8')
|
|
sendjson = json.dumps({'type': 'CMD', 'source': self.network.own_addr, 'nonce': nonce, 'message': ct}).encode('UTF-8')
|
|
self.network.send_msg(self.serverAddr,sendjson)
|
|
|
|
def recieveMessage(self, message: bytes) -> bytes:
|
|
status, msg = self.network.receive_msg(blocking=True)
|
|
if not status:
|
|
raise Exception('Network error during connection.')
|
|
try:
|
|
b64 = json.loads(msg)
|
|
retnonce = b64decode(b64['nonce'])
|
|
retciphertext = b64decode(b64['message'])
|
|
retcipher = ChaCha20.new(self.cipherkey, nonce=retnonce)
|
|
plaintext = retcipher.decrypt(retciphertext)
|
|
return plaintext
|
|
except Exception:
|
|
print("Incorrect decryption")
|