Implement reference counting mechanism

tfw/internals/crypto/__init__.py

@@ -0,0 +1,2 @@
+from .authentication import sign_message, verify_message
+from .key_manager import KeyManager

tfw/internals/crypto/authentication.py

@@ -0,0 +1,71 @@
+from functools import wraps
+from base64 import b64encode, b64decode
+from copy import deepcopy
+
+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
+
+
+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 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

tfw/internals/crypto/key_manager.py

@@ -0,0 +1,48 @@
+from atexit import register
+from tempfile import gettempdir
+from os import urandom, chmod, remove
+from os.path import exists, join
+from stat import S_IRUSR, S_IWUSR, S_IXUSR
+
+from tfw.internals.lazy import lazy_property
+from tfw.internals.ref_counter import RefCounter
+
+
+KEYFILE = join(gettempdir(), 'tfw-auth.key')
+LOCKFILE = join(gettempdir(), 'tfw-auth.lock')
+
+
+class KeyManagerRefCounter(RefCounter):
+    def deallocate(self):
+        if exists(KEYFILE):
+            remove(KEYFILE)
+
+
+class KeyManager:
+    keyfile = KEYFILE
+    refcounter = KeyManagerRefCounter(LOCKFILE)
+
+    def __init__(self):
+        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)
+
+register(KeyManager.refcounter.teardown_instance)