baseimage-tutorial-framework/tfw/internals/crypto/key_manager.py

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)