from functools import wraps from base64 import b64encode, b64decode from copy import deepcopy from hashlib import md5 from os import urandom, chmod from os.path import exists from stat import S_IRUSR, S_IWUSR, S_IXUSR from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives.hashes import SHA256 from cryptography.hazmat.primitives.hmac import HMAC as _HMAC from cryptography.exceptions import InvalidSignature from tfw.networking import message_bytes from tfw.decorators.lazy_property import lazy_property from tfw.config import TFWENV def message_checksum(message): return md5(message_bytes(message)).hexdigest() def sign_message(key, message): signature = message_signature(key, message) message['signature'] = b64encode(signature).decode() def message_signature(key, message): return HMAC(key, message_bytes(message)).signature def verify_message(key, message): message = deepcopy(message) try: signature_b64 = message.pop('signature') signature = b64decode(signature_b64) actual_signature = message_signature(key, message) return signature == actual_signature except KeyError: return False class KeyManager: def __init__(self): self.keyfile = TFWENV.AUTH_KEY if not exists(self.keyfile): self._init_auth_key() @lazy_property def auth_key(self): with open(self.keyfile, 'rb') as ifile: return ifile.read() def _init_auth_key(self): key = self.generate_key() with open(self.keyfile, 'wb') as ofile: ofile.write(key) self._chmod_700_keyfile() return key @staticmethod def generate_key(): return urandom(32) def _chmod_700_keyfile(self): chmod(self.keyfile, S_IRUSR | S_IWUSR | S_IXUSR) class HMAC: def __init__(self, key, message): self.key = key self.message = message self._hmac = _HMAC( key=key, algorithm=SHA256(), backend=default_backend() ) def _reload_if_finalized(f): # pylint: disable=no-self-argument,not-callable @wraps(f) def wrapped(instance, *args, **kwargs): if getattr(instance, '_finalized', False): instance.__init__(instance.key, instance.message) ret_val = f(instance, *args, **kwargs) setattr(instance, '_finalized', True) return ret_val return wrapped @property @_reload_if_finalized def signature(self): self._hmac.update(self.message) signature = self._hmac.finalize() return signature @_reload_if_finalized def verify(self, signature): self._hmac.update(self.message) try: self._hmac.verify(signature) return True except InvalidSignature: return False