baseimage-tutorial-framework/lib/tfw/crypto.py

105 lines
2.8 KiB
Python
Raw Normal View History

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
2019-05-27 12:09:13 +00:00
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