mirror of
https://github.com/avatao-content/baseimage-tutorial-framework
synced 2024-11-14 04:47:17 +00:00
108 lines
2.9 KiB
Python
108 lines
2.9 KiB
Python
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.internals.networking import message_bytes
|
|
from tfw.internals.lazy import lazy_property
|
|
from tfw.config import TFWENV
|
|
|
|
|
|
def message_checksum(message):
|
|
return md5(message_bytes(message)).hexdigest()
|
|
|
|
|
|
def sign_message(key, message):
|
|
message.pop('scope', None)
|
|
message.pop('signature', None)
|
|
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.pop('scope', None)
|
|
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
|