DecentraSanta/classes/SantasBrain.py

import base64
import hashlib
import json
import logging
import math
import random
import secrets
import threading
from collections import defaultdict
from typing import Optional, Callable

import Crypto.Util.number
from Crypto.Cipher import AES
from Crypto.Hash import SHAKE128
from Crypto.Protocol.DH import key_agreement
from Crypto.PublicKey import ECC
from Crypto.PublicKey.ECC import EccKey
from Crypto.Util.number import getPrime

from classes import Message
from classes.Crypto.CSPRNG import CSPRNG
from classes.Crypto.CommutativeCipher import CommutativeCipher
from classes.MessageHandler import MessageHandler
from classes.MessageTypes.Announcement import AnnouncementMessage
from classes.MessageTypes.Introduction import IntroductionMessage
from classes.MessageTypes.Ready import ReadyMessage
from classes.MessageTypes.Shuffle import ShuffleMessage
from classes.UserInterface import UserInterface

logger = logging.getLogger(__name__)

SHUFFLE_CARDS_STAGE = 'shuffle_cards'
DECRYPT_CARDS_STAGE = 'decrypt_cards'
BUILD_ANONYMOUS_STAGE = 'build_announcement'
ENCRYPT_ANONYMOUS_STAGE = 'encrypt_announcement'
SHUFFLE_ANONYMOUS_STAGE = 'shuffle_announcement'
DECRYPT_ANONYMOUS_STAGE = 'decrypt_announcement'

class Brain:
	def __init__(self, message_handler: MessageHandler, user_interface: UserInterface):
		self.thread_lock = threading.Lock()
		self.message_handler = message_handler
		self.user_interface = user_interface

		# Store the names, public key and seed commits of each participant. We will assume that each participant uses
		# a unique name. It is expected that the participants in the secret santa will make sure that this condition is
		# met. If a non-unique name is received, throw an error to warn the user, and then ignore the new participant
		# with the existing name.
		self.known_participants: dict[str, tuple[EccKey, bytes]] = {}
		# Store who we expect to take part in the secret santa
		self.chosen_participants: Optional[list[str]] = None
		# Store all received messages (except Introduction messages), by sender
		self.received_messages: dict[str, list[Message]] = defaultdict(list)
		# Our own info, that we will send out
		self.own_name: Optional[str] = None
		self.signing_key: Optional[EccKey] = None
		self.random_seed: Optional[bytes] = None
		self.introduction_message: Optional[IntroductionMessage] = None
		self.info_for_santa: Optional[str] = None

		# Secret key, used for Diffie-Hellman exchange later
		self.secret_key = ECC.generate(curve='p256')

		self.user_interface.add_user_info_listener(self.receive_user_info)
		self.user_interface.add_start_listener(self.receive_user_start_command)
		self.message_handler.add_message_receiver(self.receive_message)

		# The variables we need through the secret santa process itself
		self.process_failed = False
		self.card_values: Optional[list[bytes]] = None
		self.card_exchange_cipher: Optional[CommutativeCipher] = None
		self.announcement_build_cipher: Optional[CommutativeCipher] = None
		self.announcement_shuffle_cipher: Optional[CommutativeCipher] = None
		self.sent_card_shuffling: bool = False
		self.sent_card_decryption: bool = False
		self.card_drawn: Optional[int] = None
		self.built_anonymous_deck = False
		self.encrypted_anonymous_deck = False
		self.shuffled_anonymous_deck = False
		self.decrypted_anonymous_deck = False
		self.anonymous_keys: Optional[dict[int, EccKey]] = None
		self.sent_announcement = False
		self.received_announcement = False

	def receive_user_info(self, name: str, info_for_santa: str):
		# We will only receive this once
		if (
			self.own_name is not None
			or self.signing_key is not None
			or self.introduction_message is not None
			or self.info_for_santa is not None
			or self.random_seed is not None
		):
			return

		self.own_name = name
		self.signing_key = ECC.generate(curve='p256')
		self.random_seed = secrets.token_bytes(32)
		self.info_for_santa = info_for_santa

		# Need to create a commit for our random seed
		hasher = hashlib.sha512()
		hasher.update(self.random_seed)
		seed_commit = hasher.digest()
		self.introduction_message = IntroductionMessage(name, self.signing_key.public_key(), seed_commit)
		self.message_handler.send_message(self.introduction_message, self.signing_key)

	def receive_user_start_command(self, chosen_participants: list[str]):
		with self.thread_lock:
			# If we have not introduced ourselves to the world and selected a random seed and all that, ignore command
			if self.own_name is None or self.random_seed is None:
				return
			# If we have already selected who we want to exchange with, ignore this
			if self.chosen_participants is not None:
				return

			# Make sure that each participant we want to play with actually exists
			confirmed_existing = []

			# Make sure that no participant is repeated in the list
			chosen_participants = list(set(chosen_participants))

			for name in chosen_participants:
				if name not in self.known_participants:
					logger.error(f'Tried to start a secret santa with non-existing participant {name}')
					return
				key, _ = self.known_participants[name]
				confirmed_existing.append((name, key))

			ready_message = ReadyMessage(self.own_name, confirmed_existing, self.random_seed)
			self.chosen_participants = sorted([name for name, key in confirmed_existing])
		self.message_handler.send_message(ready_message, self.signing_key)

	def receive_message(self, message: Message):
		# If this is an introduction message, it needs special handling
		if isinstance(message, IntroductionMessage):
			self.receive_introduction_message(message)
		# For normal, non-introduction messages, check that they're from who they purport to be from,
		# and add them to the list
		else:
			with self.thread_lock:
				name = message.get_name()
				if name not in self.known_participants:
					logger.warning(f'Received message from unknown participant {name}. Ignoring.')
					return
				key, _ = self.known_participants[name]
				if not message.check_signature(key):
					logger.warning(f'Received message that purports to be from {name} with invalid signature. Ignoring.')
					return
				self.received_messages[name].append(message)

			# Each received message triggers the main function of this class
			messages_to_send = self.santa_loop()
			for message in messages_to_send:
				self.message_handler.send_message(message, self.signing_key)

	def receive_introduction_message(self, message: IntroductionMessage):
		discovered_new_participant = False
		name = message.get_name()
		key = message.get_key()

		with self.thread_lock:
			if name in self.known_participants and self.known_participants[name][0] != key:
				logger.error(
					f'There are two participants using the name {name}. All participants need unique name. '
					f'The second participant will be ignored, but this may lead to the santa exchange not being started.'
				)
				return

			if name not in self.known_participants:
				discovered_new_participant = True
				self.known_participants[name] = (key, message.get_seed_commit())

		# If this a participant we don't already know about, chances are they don't know about us, either.
		# Send an introduction message, if we are ready for that
		if (
				discovered_new_participant
				and self.introduction_message is not None
				and self.signing_key is not None
		):
			self.user_interface.receive_user(name)
			self.message_handler.send_message(self.introduction_message, self.signing_key)


	# Santa's brain will be driven by receiving messages. We'll call this main method each time we receive a message.
	# This is probably inefficient, but it makes it easier to follow what the code is doing
	def santa_loop(self) -> list[Message]:
		# We don't want to send messages while holding the thread lock, in case this leads to a deadlock
		messages_to_send: list[Message] = []
		with self.thread_lock:
			# If something has caused the process to fail, don't try to do it again
			if self.process_failed:
				return messages_to_send
			# First of all, if the user hasn't pressed start yet, there is nothing to do here
			if self.chosen_participants is None:
				return messages_to_send
			# Next, check that the user has actually provided all the information we need to perform this process
			if self.own_name is None or self.info_for_santa is None or self.random_seed is None:
				return messages_to_send

			# Next, if we haven't built our commutative ciphers and card values yet, attempt to do that first
			if (
				self.card_values is None
				or self.card_exchange_cipher is None
				or self.announcement_build_cipher is None
				or self.announcement_shuffle_cipher is None
			):
				should_continue = self.build_ciphers()
				if not should_continue:
					return messages_to_send

			# We will give each participant a number, based on their alphabetical order
			all_participants = sorted(list(set(self.chosen_participants + [self.own_name])))

			# Shuffle cards
			if not self.sent_card_shuffling:
				shuffle_message = self.build_shuffle_message(all_participants)
				if shuffle_message is None:
					return messages_to_send
				messages_to_send.append(shuffle_message)
				self.sent_card_shuffling = True

			# Decrypt shuffled cards
			if not self.sent_card_decryption:
				decrypt_cards_message = self.build_decrypt_cards_message(all_participants)
				if decrypt_cards_message is None:
					return messages_to_send
				messages_to_send.append(decrypt_cards_message)
				self.sent_card_decryption = True

			# Find which card we drew
			if self.card_drawn is None:
				own_index = all_participants.index(self.own_name)
				last_participant = all_participants[-1]
				message = next(
					(
						message for message in self.received_messages[last_participant]
						if isinstance(message, ShuffleMessage) and message.get_stage() == DECRYPT_CARDS_STAGE
					),
					None
				)
				if message is None:
					return messages_to_send
				self.decrypt_card_value(message.get_cards()[own_index])

			# Next, anonymously publish a key someone else can use to encrypt a message telling you that you're
			# their secret santa
			if not self.built_anonymous_deck:
				build_deck_message = self.build_anonymous_deck(all_participants)
				if build_deck_message is None:
					return messages_to_send
				messages_to_send.append(build_deck_message)
				self.built_anonymous_deck = True

			if not self.encrypted_anonymous_deck:
				encrypt_deck_message = self.encrypt_anonymous_deck(all_participants)
				if encrypt_deck_message is None:
					return messages_to_send
				messages_to_send.append(encrypt_deck_message)
				self.encrypted_anonymous_deck = True

			if not self.shuffled_anonymous_deck:
				shuffle_anonymous_deck_message = self.shuffle_anonymous_deck(all_participants)
				if shuffle_anonymous_deck_message is None:
					return messages_to_send
				messages_to_send.append(shuffle_anonymous_deck_message)
				self.shuffled_anonymous_deck = True

			if not self.decrypted_anonymous_deck:
				decrypted_anonymous_deck_message = self.decrypt_anonymous_deck(all_participants)
				if decrypted_anonymous_deck_message is None:
					return messages_to_send
				messages_to_send.append(decrypted_anonymous_deck_message)
				self.decrypted_anonymous_deck = True

			if not self.sent_announcement:
				announcement_message = self.build_announcement(all_participants)
				if announcement_message is None:
					return messages_to_send
				messages_to_send.append(announcement_message)
				self.sent_announcement = True

			# Look through all announcements to find our secret santa receiver
			if self.sent_announcement and self.anonymous_keys is not None and not self.received_announcement:
				receiver_card_no = (self.card_drawn + 1) % len(all_participants)
				announcements = []
				for participant in self.received_messages:
					for message in self.received_messages[participant]:
						if isinstance(message, AnnouncementMessage):
							announcements.append(message)

				receiver_key = self.anonymous_keys[receiver_card_no]
				def kdf(x):
					return SHAKE128.new(x).read(32)
				session_key = key_agreement(eph_priv=self.secret_key, eph_pub=receiver_key, kdf=kdf)

				for announcement in announcements:
					announcement_hash = announcement.get_announcement_hash()
					encrypted_announcement = announcement.get_encrypted_announcement().decode('utf-8')
					try:
						encrypted_announcement = json.loads(encrypted_announcement)
						ciphertext = base64.b64decode(encrypted_announcement['ciphertext'])
						tag = base64.b64decode(encrypted_announcement['tag'])
						nonce = base64.b64decode(encrypted_announcement['nonce'])

						cipher = AES.new(session_key, AES.MODE_EAX, nonce=nonce)
						plaintext = cipher.decrypt(ciphertext)
						cipher.verify(tag)

						plaintext = json.loads(plaintext)
						name = plaintext['name']
						extra = plaintext['extra']

						self.user_interface.announce_recipient(name, extra)
						self.received_announcement = True
					except:
						continue
		return messages_to_send


	def build_ciphers(self) -> bool:
		received_seeds = [self.random_seed]
		for name in self.chosen_participants:
			_, seed_commit = self.known_participants[name]
			ready_message = next(
				(message for message in self.received_messages[name] if isinstance(message, ReadyMessage)),
				None
			)
			if ready_message is None:
				# We haven't received ready messages from everyone yet, and so we cannot continue at the moment
				return False

			# Check that the seed we received from this user matches the seed they committed to sending
			received_seed = ready_message.get_random_seed()
			hasher = hashlib.sha512()
			hasher.update(received_seed)
			expected_commit = hasher.digest()
			if expected_commit != seed_commit:
				logger.critical(
					f'Received random seed from user {name} that did not match their commit. '
					f'Cancelling secret santa exchange!'
				)
				self.process_failed = True
				return False

			received_seeds.append(received_seed)
		# Next, create a combined random seed from all the provided seeds
		number_of_seed_bytes = max(len(seed) for seed in received_seeds)
		total_seed = b'\0' * number_of_seed_bytes
		for received_seed in received_seeds:
			# Pad seed with leading zeros if not long enough
			while len(received_seed) < number_of_seed_bytes:
				received_seed = b'\0' + received_seed
			# Xor seed with current seed so far
			total_seed = bytes(a ^ b for a, b in zip(total_seed, received_seed))

		# Use this seed in a cryptographically secure random number generator
		random_generator = CSPRNG(total_seed)

		p = Crypto.Util.number.getPrime(1500, randfunc=random_generator.get_random_bytes)
		q = Crypto.Util.number.getPrime(1000, randfunc=random_generator.get_random_bytes)

		self.card_exchange_cipher = CommutativeCipher(p, q)
		self.announcement_build_cipher = CommutativeCipher(p, q)
		self.announcement_shuffle_cipher = CommutativeCipher(p, q)

		self.card_values = [random_generator.get_random_bytes(8) for i in range(len(self.chosen_participants) + 1)]
		return True

	def build_shuffle_message(self, all_participants: list[str]) -> Optional[ShuffleMessage]:
		own_index = all_participants.index(self.own_name)
		if own_index == 0:
			# If we are the first participant, we must create the list
			card_deck = [card_value for card_value in self.card_values]
		else:
			previous_participant = all_participants[own_index - 1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == SHUFFLE_CARDS_STAGE
				), None
			)
			# If we haven't received a shuffle message yet, we can't continue
			if message is None:
				return None

			card_deck = [card for card in message.get_cards()]

		# Shuffle by drawing random numbers from secret
		shuffled_deck = []
		while len(card_deck) > 0:
			drawn_card = secrets.randbelow(len(card_deck))
			shuffled_deck.append(self.card_exchange_cipher.encode(card_deck[drawn_card]))
			del card_deck[drawn_card]

		return ShuffleMessage(self.own_name, shuffled_deck, SHUFFLE_CARDS_STAGE)

	def build_decrypt_cards_message(self, all_participants: list[str]) -> Optional[ShuffleMessage]:
		own_index = all_participants.index(self.own_name)
		# Again, special case if we are the first participant
		message = None
		previous_participant = all_participants[(own_index - 1) % len(all_participants)]
		if own_index == 0:
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == SHUFFLE_CARDS_STAGE
				),
				None
			)
		else:
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == DECRYPT_CARDS_STAGE
				),
				None
			)
		if message is None:
			return None

		card_deck = message.get_cards()
		decrypted_cards = []
		for index, card_value in enumerate(card_deck):
			# Don't decrypt our own card
			if index == own_index:
				decrypted_cards.append(card_value)
			else:
				decrypted_cards.append(self.card_exchange_cipher.decode(card_value))

		# In case we are the last participant in the list, we are actually ready to fully decrypt
		if own_index == len(all_participants) - 1:
			self.decrypt_card_value(decrypted_cards[-1])

		return ShuffleMessage(self.own_name, decrypted_cards, DECRYPT_CARDS_STAGE)

	def decrypt_card_value(self, card: bytes):
		decrypted_card_bytes = self.card_exchange_cipher.decode(card)
		if decrypted_card_bytes not in self.card_values:
			logging.critical(f'Received an invalid card after shuffling. Secret santa exchange failed!')
			self.process_failed = True
			return

		self.card_drawn = self.card_values.index(decrypted_card_bytes)

	def build_anonymous_deck(self, all_participants: list[str]) -> Optional[ShuffleMessage]:
		own_index = all_participants.index(self.own_name)

		card_deck = None
		if own_index == 0:
			card_deck = []
		else:
			previous_participant = all_participants[own_index-1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == BUILD_ANONYMOUS_STAGE
				),
				None
			)

			if message is None:
				return None

			card_deck = [card for card in message.get_cards()]

		anonymous_message = {
			'card_no': self.card_drawn,
			'key': self.secret_key.public_key().export_key(format='OpenSSH')
		}
		anonymous_message = json.dumps(anonymous_message).encode('utf-8')
		card_deck.append(self.announcement_build_cipher.encode(anonymous_message))

		return ShuffleMessage(self.own_name, card_deck, BUILD_ANONYMOUS_STAGE)

	def encrypt_anonymous_deck(self, all_participants: list[str]) -> Optional[ShuffleMessage]:
		own_index = all_participants.index(self.own_name)
		card_deck = None
		if own_index == 0:
			previous_participant = all_participants[-1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == BUILD_ANONYMOUS_STAGE
				),
				None
			)
			if message is not None:
				card_deck = message.get_cards()
		else:
			previous_participant = all_participants[own_index - 1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == ENCRYPT_ANONYMOUS_STAGE
				),
				None
			)
			if message is not None:
				card_deck = message.get_cards()

		if card_deck is None:
			return None

		encrypted_deck = []
		for index, card_value in enumerate(card_deck):
			if index == own_index:
				encrypted_deck.append(card_value)
			else:
				encrypted_deck.append(self.announcement_build_cipher.encode(card_value))

		return ShuffleMessage(self.own_name, encrypted_deck, ENCRYPT_ANONYMOUS_STAGE)

	def shuffle_anonymous_deck(self, all_participants: list[str]) -> Optional[ShuffleMessage]:
		own_index = all_participants.index(self.own_name)
		card_deck = None
		if own_index == 0:
			previous_participant = all_participants[-1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == ENCRYPT_ANONYMOUS_STAGE
				),
				None
			)
			if message is not None:
				card_deck = [card_value for card_value in message.get_cards()]
		else:
			previous_participant = all_participants[own_index-1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == SHUFFLE_ANONYMOUS_STAGE
				),
				None
			)
			if message is not None:
				card_deck = [card_value for card_value in message.get_cards()]

		if card_deck is None:
			return None

		shuffled_cards = []
		while len(card_deck) > 0:
			draw_number = secrets.randbelow(len(card_deck))
			card_value = card_deck[draw_number]
			del card_deck[draw_number]

			decrypted_previous = self.announcement_build_cipher.decode(card_value)
			shuffled_cards.append(self.announcement_shuffle_cipher.encode(decrypted_previous))

		return ShuffleMessage(self.own_name, shuffled_cards, SHUFFLE_ANONYMOUS_STAGE)


	def decrypt_anonymous_deck(self, all_participants: list[str]) -> Optional[ShuffleMessage]:
		own_index = all_participants.index(self.own_name)
		card_deck = None
		if own_index == 0:
			previous_participant = all_participants[-1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == SHUFFLE_ANONYMOUS_STAGE
				),
				None
			)

			if message is not None:
				card_deck = [card_value for card_value in message.get_cards()]
		else:
			previous_participant = all_participants[own_index - 1]
			message = next(
				(
					message for message in self.received_messages[previous_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == DECRYPT_ANONYMOUS_STAGE
				),
				None
			)

			if message is not None:
				card_deck = [card_value for card_value in message.get_cards()]

		if card_deck is None:
			return None

		decrypted_deck = [self.announcement_shuffle_cipher.decode(card_value) for card_value in card_deck]
		shuffle_message = ShuffleMessage(self.own_name, decrypted_deck, DECRYPT_ANONYMOUS_STAGE)

		# If we are the last participant, everything should now be decrypted
		if own_index == len(all_participants) - 1:
			self.get_anonymous_keys(shuffle_message)

		return shuffle_message


	def get_anonymous_keys(self, message: ShuffleMessage):
		anonymous_keys = {}
		cards = message.get_cards()
		for card in cards:
			try:
				decoded_card = json.loads(card.decode('utf-8'))
				anonymous_keys[decoded_card['card_no']] = ECC.import_key(decoded_card['key'])
			except:
				logger.critical(f'Received card {card} could not be decoded as JSON. Secret santa process failed.')
				self.process_failed = True
				return
		self.anonymous_keys = anonymous_keys

	def build_announcement(self, all_participants: list[str]) -> Optional[AnnouncementMessage]:
		if self.anonymous_keys is None:
			last_participant = all_participants[-1]
			message = next(
				(
					message for message in self.received_messages[last_participant]
					if isinstance(message, ShuffleMessage) and message.get_stage() == DECRYPT_ANONYMOUS_STAGE
				),
				None
			)

			if message is None:
				return None

			self.get_anonymous_keys(message)

		if self.anonymous_keys is None:
			return None


		santa_card = (self.card_drawn - 1) % len(all_participants)

		santa_key = self.anonymous_keys[santa_card]
		def kdf(x):
			return SHAKE128.new(x).read(32)

		session_key = key_agreement(eph_priv=self.secret_key, eph_pub=santa_key, kdf=kdf)
		session_cipher = AES.new(session_key, AES.MODE_EAX)

		message_string = {'name': self.own_name, 'extra': self.info_for_santa}

		ciphertext, tag = session_cipher.encrypt_and_digest(json.dumps(message_string).encode('utf-8'))
		hasher = hashlib.sha512()
		hasher.update(self.info_for_santa.encode('utf-8'))
		hashed_announcement = hasher.digest()

		encrypted_announcement = {
			'ciphertext': base64.b64encode(ciphertext).decode('utf-8'),
			'tag': base64.b64encode(tag).decode('utf-8'),
			'nonce': base64.b64encode(session_cipher.nonce).decode('utf-8')
		}

		encrypted_announcement = json.dumps(encrypted_announcement).encode('utf-8')
		return AnnouncementMessage(self.own_name, encrypted_announcement, hashed_announcement)