DecentraSanta/classes/MessageHandler.py

import base64
import binascii
import json
import logging
from collections.abc import Callable
from json import JSONDecodeError

from Crypto.PublicKey import ECC
from Crypto.PublicKey.ECC import EccKey

from classes.Message import Message
from classes.MessageTypes.Announcement import AnnouncementMessage
from classes.MessageTypes.Introduction import IntroductionMessage
from classes.MessageTypes.Ready import ReadyMessage
from classes.MessageTypes.Shuffle import ShuffleMessage

logger = logging.getLogger(__name__)


class MessageHandler:
	def __init__(self):
		self.receivers: list[Callable[[Message], None]] = []

	def send_message(self, message: Message, signing_key: EccKey):
		# Must be implemented by child classes
		pass

	def add_message_receiver(self, message_receiver: Callable[[Message], None]):
		self.receivers.append(message_receiver)

	def remove_message_receiver(self, message_receiver: Callable[[Message], None]):
		self.receivers.remove(message_receiver)

	def decode_received_message(self, message_string: str):
		try:
			message_object = json.loads(message_string)
		except JSONDecodeError | UnicodeDecodeError:
			logger.error(f'Could not decode received string {message_string}')
			return

		message_type = message_object.get('type')
		if message_type is None:
			logger.error(f'Message type not found in message {message_string}')
			return

		message = None

		if message_type == 'introduction':
			name = message_object.get('name')
			seed_commit = message_object.get('seed_commit')
			key = message_object.get('key')

			if None in [name, seed_commit, key]:
				logger.error(f'Did not find expected fields for introduction message: {message_string}')
				return

			elif not isinstance(name, str) or not isinstance(seed_commit, str) or not isinstance(key, str):
				logger.error(f'Received data of the wrong type for introduction message: {message_string}')
				return

			try:
				key = ECC.import_key(key)
			except ValueError:
				logger.error(f'{key} is not a valid key')
				return

			try:
				seed_commit = base64.b64decode(seed_commit, validate=True)
			except binascii.Error:
				logger.error(f'Seed commit {seed_commit} is not a valid base64 string')
				return

			message = IntroductionMessage(name, key, seed_commit)

		elif message_type == 'ready':
			name = message_object.get('name')
			participants = message_object.get('participants')
			random_seed = message_object.get('random_seed')

			if None in [name, participants, random_seed]:
				logger.error(f'Did not find expected fields for ready message {message_string}')
				return

			elif not isinstance(name, str) or not isinstance(participants, list) or not isinstance(random_seed, str):
				logger.error(f'Received data of the wrong type for ready message: {message_string}')
				return

			elif not all(isinstance(participant, list) for participant in participants):
				logger.error(f'Not all participants in participant list are lists: {message_string}')
				return

			elif not all(len(participant_tuple) == 2 for participant_tuple in participants):
				logger.error(f'Not all participant tuples are of length two in {message_string}')
				return

			elif not all(isinstance(name, str) and isinstance(key, str) for name, key in participants):
				logger.error(f'Not all participant tuples contain a name and a key')
				return

			decoded_participants = []
			for participant_name, key in participants:
				try:
					key = ECC.import_key(key)
					decoded_participants.append((participant_name, key))
				except ValueError:
					logger.error(f'Could not decode public key {key} of participant {participant_name}')
					return

			try:
				random_seed = base64.b64decode(random_seed, validate=True)
			except binascii.Error:
				logger.error(f'Random seed {random_seed} from {name} is not a valid base64 string')
				return

			message = ReadyMessage(name, decoded_participants, random_seed)

		elif message_type == 'shuffle':
			name = message_object.get('name')
			cards = message_object.get('cards')
			stage = message_object.get('stage')

			if None in [name, cards, stage]:
				logger.error(f'Did not receive all expected fields for shuffle message: {message_string}')
				return

			elif not isinstance(name, str) or not isinstance(cards, list) or not isinstance(stage, str):
				logger.error(f'Received fields were not correct type for shuffle message: {message_string}')
				return

			elif not all(isinstance(card, str) for card in cards):
				logger.error(f'All received cards were not of type string: {message_string}')
				return

			new_cards = []
			for card in cards:
				try:
					new_cards.append(base64.b64decode(card, validate=True))
				except binascii.Error:
					logger.error(f'{card} is not a valid base64 string')
					return

			message = ShuffleMessage(name, new_cards, stage)

		elif message_type == 'announcement':
			name = message_object.get('name')
			announcement = message_object.get('encrypted_announcement')
			announcement_hash = message_object.get('announcement_hash')

			if None in [name, announcement, announcement_hash]:
				logger.error(f'Did not receive all expected fields for announcement message: {message_string}')
				return

			elif not isinstance(name, str) or not isinstance(announcement, str) or not isinstance(announcement_hash, str):
				logger.error(f'Received fields for announcement message are not of correct type: {message_string}')
				return

			try:
				announcement = base64.b64decode(announcement, validate=True)
			except binascii.Error:
				logger.error(f'{announcement} is not a valid base64 string')
				return

			try:
				announcement_hash = base64.b64decode(announcement_hash, validate=True)
			except binascii.Error:
				logger.error(f'{announcement} is not a valid base64 string')
				return

			message = AnnouncementMessage(name, announcement, announcement_hash)

		if message is None:
			logger.error(f'Message type {message_type} does not exist')
			return

		if message_object.get('signature'):
			message.set_signature(base64.b64decode(message_object.get('signature')))
		for message_receiver in self.receivers:
			message_receiver(message)