openstack/deb-python-sqlalchemy-utils
Code Issues Proposed changes

Allow key to be a callable

Browse Source
This commit is contained in:
Ryan Leckey
2014-10-22 06:32:27 -07:00
parent e465616477
commit a814493a77
1 changed files with 14 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
sqlalchemy_utils/types/encrypted.py
View File

@@ -24,13 +24,14 @@ class EncryptionDecryptionBaseEngine(object):
new engines. new engines.
""" """
def __init__(self, key): def _update_key(self, key):
"""Initialize a base engine."""
if isinstance(key, six.string_types): if isinstance(key, six.string_types):
key = six.b(key) key = six.b(key)
self._digest = hashes.Hash(hashes.SHA256(), backend=default_backend()) digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
self._digest.update(key) digest.update(key)
self._engine_key = self._digest.finalize() engine_key = digest.finalize()
self._initialize_engine(engine_key)
def encrypt(self, value): def encrypt(self, value):
raise NotImplementedError('Subclasses must implement this!') raise NotImplementedError('Subclasses must implement this!')
@@ -45,14 +46,6 @@ class AesEngine(EncryptionDecryptionBaseEngine):
BLOCK_SIZE = 16 BLOCK_SIZE = 16
PADDING = six.b('*') PADDING = six.b('*')
def __init__(self, key):
super(AesEngine, self).__init__(key)
self._initialize_engine(self._engine_key)
def _update_key(self, new_key):
parent = EncryptionDecryptionBaseEngine(new_key)
self._initialize_engine(parent._engine_key)
def _initialize_engine(self, parent_class_key): def _initialize_engine(self, parent_class_key):
self.secret_key = parent_class_key self.secret_key = parent_class_key
self.iv = self.secret_key[:16] self.iv = self.secret_key[:16]
@@ -96,14 +89,6 @@ class AesEngine(EncryptionDecryptionBaseEngine):
class FernetEngine(EncryptionDecryptionBaseEngine): class FernetEngine(EncryptionDecryptionBaseEngine):
"""Provide Fernet encryption and decryption methods.""" """Provide Fernet encryption and decryption methods."""
def __init__(self, key):
super(FernetEngine, self).__init__(key)
self._initialize_engine(self._engine_key)
def _update_key(self, new_key):
parent = EncryptionDecryptionBaseEngine(new_key)
self._initialize_engine(parent._engine_key)
def _initialize_engine(self, parent_class_key): def _initialize_engine(self, parent_class_key):
self.secret_key = base64.urlsafe_b64encode(parent_class_key) self.secret_key = base64.urlsafe_b64encode(parent_class_key)
self.fernet = Fernet(self.secret_key) self.fernet = Fernet(self.secret_key)
@@ -212,7 +197,7 @@ class EncryptedType(TypeDecorator):
self._key = key self._key = key
if not engine: if not engine:
engine = AesEngine engine = AesEngine
self.engine = engine(self._key) self.engine = engine()
@property @property
def key(self): def key(self):
@@ -221,15 +206,20 @@ class EncryptedType(TypeDecorator):
@key.setter @key.setter
def key(self, value): def key(self, value):
self._key = value self._key = value
self.engine._update_key(self._key)
def _update_key(self):
key = self._key() if callable(self._key) else self._key
self.engine._update_key(key)
def process_bind_param(self, value, dialect): def process_bind_param(self, value, dialect):
"""Encrypt a value on the way in.""" """Encrypt a value on the way in."""
if value is not None: if value is not None:
self._update_key()
return self.engine.encrypt(value) return self.engine.encrypt(value)
def process_result_value(self, value, dialect): def process_result_value(self, value, dialect):
"""Decrypt value on the way out.""" """Decrypt value on the way out."""
if value is not None: if value is not None:
self._update_key()
decrypted_value = self.engine.decrypt(value) decrypted_value = self.engine.decrypt(value)
return self.underlying_type.python_type(decrypted_value) return self.underlying_type.python_type(decrypted_value)