import base64
import json
from Crypto.Hash import SHA256
from Crypto.PublicKey.ECC import EccKey
from Crypto.Signature import DSS

class Message:
	def __init__(self, fields = None):
		self.message_fields = {}
		if fields is not None:
			for field in fields:
				self.message_fields[field] = fields[field]

	def generate_and_sign(self, private_key: EccKey):
		if 'signature' in self.message_fields:
			del self.message_fields['signature']

		message = json.dumps(self.message_fields)
		h = SHA256.new(message.encode('utf-8'))
		signer = DSS.new(private_key, 'fips-186-3')
		signature = signer.sign(h)
		self.message_fields['signature'] = base64.b64encode(signature).decode('utf-8')
		return json.dumps(self.message_fields)

	def check_signature(self, public_key: EccKey):
		signature = base64.b64decode(self.message_fields['signature'])
		message_copy = self.message_fields.copy()
		del message_copy['signature']
		message = json.dumps(message_copy)
		h = SHA256.new(message.encode('utf-8'))
		verifier = DSS.new(public_key, 'fips-186-3')

		try:
			verifier.verify(h, signature)
			return True
		except ValueError:
			return False

	def set_signature(self, signature: bytes):
		self.message_fields['signature'] = base64.b64encode(signature).decode('utf-8')

	def get_name(self) -> str:
		# All subclasses have a get_name function, which tells who sent the message
		pass