Public-Key Encryption Exchange and Encryption Algorithm Example

public-key.py

from math import sqrt
import base64
from itertools import count, islice
import random
from typing import List
from dataclasses import dataclass
from Crypto.Cipher import AES
import hashlib
from Crypto.Random import get_random_bytes

def is_prime(n):
    return n > 1 and all(n%i for i in islice(count(2), int(sqrt(n)-1)))

nonce = get_random_bytes(15)

# generates prime numbers between 1 and 10000
primes = [i for i in range(1,10000) if is_prime(i)]

@dataclass
class PublicKey:
	g: int
	p: int
	value: int

@dataclass
class PrivateKey:
	g: int
	p: int
	value: int


class AESCipher:

	@staticmethod
	def encrypt(message: str, key: str):
		hex_key = hashlib.sha256(key.encode()).hexdigest()[:16].encode('utf-8')

		obj = AES.new(hex_key, AES.MODE_OCB, nonce=nonce)
		return obj.encrypt_and_digest(message.encode())

	@staticmethod
	def decrypt(ciphertext, mac, key):
		hex_key = hashlib.sha256(key.encode()).hexdigest()[:16].encode('utf-8')
		cipher = AES.new(hex_key, AES.MODE_OCB, nonce=nonce)
		return cipher.decrypt_and_verify(ciphertext, mac).decode()


class User:

	def __init__(self, name:str, private_value):
		self.private_value: int = private_value
		print(f"\n{name}'s private value is {private_value}.")
		self.name = name
		self.public_keys: Dict[str, PublicKey] = {}
		self.p = self.generate_p()
		self.g = self.generate_g()
		self.compute_public_key_value()

		self.shared_secrets: Dict[str, int] = {}


	@property
	def public_key(self)-> PublicKey:
		return self.public_keys[self.name]

	def generate_p(self, lower_bound: int = 1000, upper_bound: int = 10000):
		filtered_primes: List[int] = list(filter(lambda prime: lower_bound <= prime <= upper_bound, primes))
		p = random.choice(filtered_primes)
		print(f"{self.name} generated {p} for p.")
		return p

	def generate_g(self, lower_bound: int = 2, upper_bound: int = 50):
		filtered_primes: List[int] = list(filter(lambda prime: lower_bound <= prime <= upper_bound, primes))
		while True:
			g = random.choice(filtered_primes)
			if self.p and g < self.p:
				print(f"{self.name} generated {g} for g.")
				return g

	def compute_public_key_value(self, save_state: bool = True)-> int:
		public_key_value = (self.g ** self.private_value) % self.p
		if save_state:
			self.public_keys[self.name] = PublicKey(self.g, self.p, public_key_value)
		return public_key_value

	def receive_key(self, value, party_name: str)-> None:
		print(f"{self.name} received {party_name}'s public key: {value}")
		self.public_keys[party_name] = value

		self.compute_shared_secret(party_name)

	def compute_shared_secret(self, name: str)-> int:
		other_public_key: PublicKey = self.public_keys[name]
		shared_secret = (other_public_key.value ** self.private_value) % other_public_key.p
		self.shared_secrets[name] = shared_secret

	def encrypt_message(self, receiver: str, message: str):
		public_value = self.public_keys[receiver].g ** self.private_value % self.public_keys[receiver].p
		K = self.public_keys[receiver].value ** self.private_value % self.public_keys[receiver].p
		ciphertext, mac = AESCipher.encrypt(message, key=str(K))
		return ciphertext, mac, public_value

	def send_message(self, receiver: str, message: str):
		print(f"{self.name} wants to send a message of '{message}' to {receiver}")
		ciphertext, mac, public_value = self.encrypt_message(receiver, message)

		return {
			"sender": self.name,
			"ciphertext": ciphertext,
			"mac": mac,
			"public_value": public_value
		}

	def receive_message(self, sender: str, ciphertext, mac, public_value):
		print(f"{self.name} received a message: {ciphertext} from {sender}")
		decryption_key = (public_value ** self.private_value) % self.p
		decrypted_message: str = AESCipher.decrypt(ciphertext, mac, key=str(decryption_key))
		print(f"After decrypting message, plain text is '{decrypted_message}'")
		return decrypted_message


if __name__ == "__main__":

	alice = User(name="Alice", private_value=92)
	bob = User(name="Bob", private_value=15)

	# exchange public keys
	alice.receive_key(bob.public_key, bob.name)
	bob.receive_key(alice.public_key, alice.name)

	message = bob.send_message(receiver="Alice", message="Hello, Alice.")
	alice.receive_message(**message)
  
  '''
  output of python public-key.py:

  Alice's private value is 92.
  Alice generated 1103 for p.
  Alice generated 11 for g.

  Bob's private value is 15.
  Bob generated 2699 for p.
  Bob generated 5 for g.
  Alice received Bob's public key: PublicKey(g=5, p=2699, value=1319)
  Bob received Alice's public key: PublicKey(g=11, p=1103, value=564)
  Bob wants to send a message of 'Hello, Alice.' to Alice
  Alice received a message: b"\xb3\x17\x02\xad1'\xe7!\xab&\xfb\xed\x0b" from Bob
  After decrypting message, plain text is 'Hello, Alice.
  '''