Slight code cleanup

classes/SantasBrain.py

@@ -2,12 +2,10 @@ 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
+from typing import Optional
 
 import Crypto.Util.number
 from Crypto.Cipher import AES
@@ -15,7 +13,6 @@ 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
@@ -36,6 +33,7 @@ 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()
@@ -85,11 +83,11 @@ class Brain:
 	def receive_user_info(self, name: str, info_for_santa: str):
 		# We will only receive this once
 		if (
-			self.own_name is not None
+				self.own_name is not None
-			or self.signing_key is not None
+				or self.signing_key is not None
-			or self.introduction_message is not None
+				or self.introduction_message is not None
-			or self.info_for_santa is not None
+				or self.info_for_santa is not None
-			or self.random_seed is not None
+				or self.random_seed is not None
 		):
 			return
 
@@ -145,7 +143,8 @@ class Brain:
 					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.')
+					logger.warning(
+						f'Received message that purports to be from {name} with invalid signature. Ignoring.')
 					return
 				self.received_messages[name].append(message)
 
@@ -181,7 +180,6 @@ class Brain:
 			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]:
@@ -200,10 +198,10 @@ class Brain:
 
 			# Next, if we haven't built our commutative ciphers and card values yet, attempt to do that first
 			if (
-				self.card_values is None
+					self.card_values is None
-				or self.card_exchange_cipher is None
+					or self.card_exchange_cipher is None
-				or self.announcement_build_cipher is None
+					or self.announcement_build_cipher is None
-				or self.announcement_shuffle_cipher is None
+					or self.announcement_shuffle_cipher is None
 			):
 				should_continue = self.build_ciphers()
 				if not should_continue:
@@ -290,12 +288,13 @@ class Brain:
 							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)
@@ -317,7 +316,6 @@ class Brain:
 						continue
 		return messages_to_send
 
-
 	def build_ciphers(self) -> bool:
 		received_seeds = [self.random_seed]
 		for name in self.chosen_participants:
@@ -364,7 +362,7 @@ class Brain:
 		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)]
+		self.card_values = [random_generator.get_random_bytes(8) for _ in range(len(self.chosen_participants) + 1)]
 		return True
 
 	def build_shuffle_message(self, all_participants: list[str]) -> Optional[ShuffleMessage]:
@@ -398,7 +396,6 @@ class Brain:
 	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(
@@ -446,11 +443,10 @@ class Brain:
 	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]
+			previous_participant = all_participants[own_index - 1]
 			message = next(
 				(
 					message for message in self.received_messages[previous_participant]
@@ -526,7 +522,7 @@ class Brain:
 			if message is not None:
 				card_deck = [card_value for card_value in message.get_cards()]
 		else:
-			previous_participant = all_participants[own_index-1]
+			previous_participant = all_participants[own_index - 1]
 			message = next(
 				(
 					message for message in self.received_messages[previous_participant]
@@ -551,7 +547,6 @@ class Brain:
 
 		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
@@ -592,7 +587,6 @@ class Brain:
 
 		return shuffle_message
 
-
 	def get_anonymous_keys(self, message: ShuffleMessage):
 		anonymous_keys = {}
 		cards = message.get_cards()
@@ -625,10 +619,10 @@ class Brain:
 		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)