44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
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 |