360 lines
14 KiB
Python
360 lines
14 KiB
Python
import hashlib
|
|
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.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.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'
|
|
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 receiving information about who your secret santa is
|
|
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
|
|
|
|
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.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])
|
|
|
|
|
|
|
|
|
|
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 = 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(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) |