diff --git a/docs/data_types.rst b/docs/data_types.rst index 3d8a0ef2..df7472c4 100644 --- a/docs/data_types.rst +++ b/docs/data_types.rst @@ -74,7 +74,7 @@ EmailType EncryptedType ------------- -.. module:: sqlalchemy_utils.types.encrypted +.. module:: sqlalchemy_utils.types.encrypted.encrypted_type .. autoclass:: EncryptedType @@ -183,4 +183,3 @@ WeekDaysType .. module:: sqlalchemy_utils.types.weekdays .. autoclass:: WeekDaysType - diff --git a/sqlalchemy_utils/types/__init__.py b/sqlalchemy_utils/types/__init__.py index 29a400e4..2ed55fce 100644 --- a/sqlalchemy_utils/types/__init__.py +++ b/sqlalchemy_utils/types/__init__.py @@ -8,7 +8,7 @@ from .country import CountryType # noqa from .currency import CurrencyType # noqa from .email import EmailType # noqa -from .encrypted import EncryptedType # noqa +from .encrypted.encrypted_type import EncryptedType # noqa from .ip_address import IPAddressType # noqa from .json import JSONType # noqa from .locale import LocaleType # noqa diff --git a/sqlalchemy_utils/types/encrypted/__init__.py b/sqlalchemy_utils/types/encrypted/__init__.py new file mode 100644 index 00000000..f3b9d2e9 --- /dev/null +++ b/sqlalchemy_utils/types/encrypted/__init__.py @@ -0,0 +1 @@ +# Module for encrypted type diff --git a/sqlalchemy_utils/types/encrypted.py b/sqlalchemy_utils/types/encrypted/encrypted_type.py similarity index 74% rename from sqlalchemy_utils/types/encrypted.py rename to sqlalchemy_utils/types/encrypted/encrypted_type.py index 9cf6c077..406be3dc 100644 --- a/sqlalchemy_utils/types/encrypted.py +++ b/sqlalchemy_utils/types/encrypted/encrypted_type.py @@ -5,8 +5,9 @@ import six from sqlalchemy.types import LargeBinary, String, TypeDecorator -from ..exceptions import ImproperlyConfigured -from .scalar_coercible import ScalarCoercible +from sqlalchemy_utils.exceptions import ImproperlyConfigured +from sqlalchemy_utils.types.encrypted.padding import PADDING_MECHANISM +from sqlalchemy_utils.types.scalar_coercible import ScalarCoercible cryptography = None try: @@ -56,7 +57,6 @@ class AesEngine(EncryptionDecryptionBaseEngine): """Provide AES encryption and decryption methods.""" BLOCK_SIZE = 16 - PADDING = six.b('*') def _initialize_engine(self, parent_class_key): self.secret_key = parent_class_key @@ -67,12 +67,21 @@ def _initialize_engine(self, parent_class_key): backend=default_backend() ) - def _pad(self, value): - """Pad the message to be encrypted, if needed.""" - BS = self.BLOCK_SIZE - P = self.PADDING - padded = (value + (BS - len(value) % BS) * P) - return padded + def _set_padding_mechanism(self, padding_mechanism=None): + """Set the padding mechanism.""" + if isinstance(padding_mechanism, six.string_types): + if padding_mechanism not in PADDING_MECHANISM.keys(): + raise ImproperlyConfigured( + "There is not padding mechanism with name {}".format( + padding_mechanism + ) + ) + + if padding_mechanism is None: + padding_mechanism = 'naive' + + padding_class = PADDING_MECHANISM[padding_mechanism] + self.padding_engine = padding_class(self.BLOCK_SIZE) def encrypt(self, value): if not isinstance(value, six.string_types): @@ -80,7 +89,7 @@ def encrypt(self, value): if isinstance(value, six.text_type): value = str(value) value = value.encode() - value = self._pad(value) + value = self.padding_engine.pad(value) encryptor = self.cipher.encryptor() encrypted = encryptor.update(value) + encryptor.finalize() encrypted = base64.b64encode(encrypted) @@ -92,7 +101,7 @@ def decrypt(self, value): decryptor = self.cipher.decryptor() decrypted = base64.b64decode(value) decrypted = decryptor.update(decrypted) + decryptor.finalize() - decrypted = decrypted.rstrip(self.PADDING) + decrypted = self.padding_engine.unpad(decrypted) if not isinstance(decrypted, six.string_types): try: decrypted = decrypted.decode('utf-8') @@ -135,18 +144,38 @@ class EncryptedType(TypeDecorator, ScalarCoercible): is decrypted. EncryptedType needs Cryptography_ library in order to work. - A simple example is given below. + + When declaring a column which will be of type EncryptedType + it is better to be as precise as possible and follow the pattern + below. .. _Cryptography: https://cryptography.io/en/latest/ :: + + a_column = sa.Column(EncryptedType(sa.Unicode, + secret_key, + FernetEngine)) + + another_column = sa.Column(EncryptedType(sa.Unicode, + secret_key, + AesEngine, + 'pkcs5')) + + + A more complete example is given below. + + :: + + import sqlalchemy as sa - from sqlalchemy.ext.declarative import declarative_base from sqlalchemy import create_engine + from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker - from sqlalchemy_utils import EncryptedType + from sqlalchemy_utils import EncryptedType + from sqlalchemy_utils.types.encrypted.encrypted_type import AesEngine secret_key = 'secretkey1234' # setup @@ -154,14 +183,27 @@ class EncryptedType(TypeDecorator, ScalarCoercible): connection = engine.connect() Base = declarative_base() + class User(Base): __tablename__ = "user" id = sa.Column(sa.Integer, primary_key=True) - username = sa.Column(EncryptedType(sa.Unicode, secret_key)) - access_token = sa.Column(EncryptedType(sa.String, secret_key)) - is_active = sa.Column(EncryptedType(sa.Boolean, secret_key)) + username = sa.Column(EncryptedType(sa.Unicode, + secret_key, + AesEngine, + 'pkcs5')) + access_token = sa.Column(EncryptedType(sa.String, + secret_key, + AesEngine, + 'pkcs5')) + is_active = sa.Column(EncryptedType(sa.Boolean, + secret_key, + AesEngine, + 'zeroes')) number_of_accounts = sa.Column(EncryptedType(sa.Integer, - secret_key)) + secret_key, + AesEngine, + 'oneandzeroes')) + sa.orm.configure_mappers() Base.metadata.create_all(connection) @@ -179,15 +221,21 @@ class User(Base): num_of_accounts = 2 user = User(username=user_name, access_token=test_token, - is_active=active, accounts_num=accounts) + is_active=active, number_of_accounts=num_of_accounts) session.add(user) session.commit() - print('id: {}'.format(user.id)) - print('username: {}'.format(user.username)) - print('token: {}'.format(user.access_token)) - print('active: {}'.format(user.is_active)) - print('accounts: {}'.format(user.accounts_num)) + user_id = user.id + + session.expunge_all() + + user_instance = session.query(User).get(user_id) + + print('id: {}'.format(user_instance.id)) + print('username: {}'.format(user_instance.username)) + print('token: {}'.format(user_instance.access_token)) + print('active: {}'.format(user_instance.is_active)) + print('accounts: {}'.format(user_instance.number_of_accounts)) # teardown session.close_all() @@ -196,9 +244,11 @@ class User(Base): engine.dispose() The key parameter accepts a callable to allow for the key to change - per-row instead of be fixed for the whole table. + per-row instead of being fixed for the whole table. :: + + def get_key(): return 'dynamic-key' @@ -212,7 +262,8 @@ class User(Base): impl = LargeBinary - def __init__(self, type_in=None, key=None, engine=None, **kwargs): + def __init__(self, type_in=None, key=None, + engine=None, padding=None, **kwargs): """Initialization.""" if not cryptography: raise ImproperlyConfigured( @@ -229,6 +280,8 @@ def __init__(self, type_in=None, key=None, engine=None, **kwargs): if not engine: engine = AesEngine self.engine = engine() + if isinstance(self.engine, AesEngine): + self.engine._set_padding_mechanism(padding) @property def key(self): diff --git a/sqlalchemy_utils/types/encrypted/padding.py b/sqlalchemy_utils/types/encrypted/padding.py new file mode 100644 index 00000000..2c587d9b --- /dev/null +++ b/sqlalchemy_utils/types/encrypted/padding.py @@ -0,0 +1,124 @@ +import six + + +class Padding(object): + """Base class for padding and unpadding.""" + + def __init__(self, block_size): + self.block_size = block_size + + def pad(value): + raise NotImplementedError('Subclasses must implement this!') + + def unpad(value): + raise NotImplementedError('Subclasses must implement this!') + + +class PKCS5Padding(Padding): + """Provide PKCS5 padding and unpadding.""" + + def pad(self, value): + if not isinstance(value, six.binary_type): + value = value.encode() + padding_length = (self.block_size - len(value) % self.block_size) + padding_sequence = padding_length * six.b(chr(padding_length)) + value_with_padding = value + padding_sequence + + return value_with_padding + + def unpad(self, value): + if isinstance(value, six.binary_type): + padding_length = value[-1] + if isinstance(value, six.string_types): + padding_length = ord(value[-1]) + value_without_padding = value[0:-padding_length] + + return value_without_padding + + +class OneAndZeroesPadding(Padding): + """Provide the one and zeroes padding and unpadding. + + This mechanism pads with 0x80 followed by zero bytes. + For unpadding it strips off all trailing zero bytes and the 0x80 byte. + """ + + BYTE_80 = 0x80 + BYTE_00 = 0x00 + + def pad(self, value): + if not isinstance(value, six.binary_type): + value = value.encode() + padding_length = (self.block_size - len(value) % self.block_size) + one_part_bytes = six.b(chr(self.BYTE_80)) + zeroes_part_bytes = (padding_length - 1) * six.b(chr(self.BYTE_00)) + padding_sequence = one_part_bytes + zeroes_part_bytes + value_with_padding = value + padding_sequence + + return value_with_padding + + def unpad(self, value): + value_without_padding = value.rstrip(six.b(chr(self.BYTE_00))) + value_without_padding = value_without_padding.rstrip( + six.b(chr(self.BYTE_80))) + + return value_without_padding + + +class ZeroesPadding(Padding): + """Provide zeroes padding and unpadding. + + This mechanism pads with 0x00 except the last byte equals + to the padding length. For unpadding it reads the last byte + and strips off that many bytes. + """ + + BYTE_00 = 0x00 + + def pad(self, value): + if not isinstance(value, six.binary_type): + value = value.encode() + padding_length = (self.block_size - len(value) % self.block_size) + zeroes_part_bytes = (padding_length - 1) * six.b(chr(self.BYTE_00)) + last_part_bytes = six.b(chr(padding_length)) + padding_sequence = zeroes_part_bytes + last_part_bytes + value_with_padding = value + padding_sequence + + return value_with_padding + + def unpad(self, value): + if isinstance(value, six.binary_type): + padding_length = value[-1] + if isinstance(value, six.string_types): + padding_length = ord(value[-1]) + value_without_padding = value[0:-padding_length] + + return value_without_padding + + +class NaivePadding(Padding): + """Naive padding and unpadding using '*'. + + The class is provided only for backwards compatibility. + """ + + CHARACTER = six.b('*') + + def pad(self, value): + num_of_bytes = (self.block_size - len(value) % self.block_size) + value_with_padding = value + num_of_bytes * self.CHARACTER + + return value_with_padding + + def unpad(self, value): + value_without_padding = value.rstrip(self.CHARACTER) + + return value_without_padding + + +PADDING_MECHANISM = { + 'pkcs5': PKCS5Padding, + 'oneandzeroes': OneAndZeroesPadding, + 'zeroes': ZeroesPadding, + 'naive': NaivePadding +} diff --git a/tests/types/test_encrypted.py b/tests/types/test_encrypted.py index fff6bfa9..d4c436fd 100644 --- a/tests/types/test_encrypted.py +++ b/tests/types/test_encrypted.py @@ -4,7 +4,7 @@ import sqlalchemy as sa from sqlalchemy_utils import ColorType, EncryptedType, PhoneNumberType -from sqlalchemy_utils.types.encrypted import ( +from sqlalchemy_utils.types.encrypted.encrypted_type import ( AesEngine, DatetimeHandler, FernetEngine @@ -18,7 +18,7 @@ @pytest.fixture -def User(Base, encryption_engine, test_key): +def User(Base, encryption_engine, test_key, padding_mechanism): class User(Base): __tablename__ = 'user' id = sa.Column(sa.Integer, primary_key=True) @@ -26,61 +26,71 @@ class User(Base): username = sa.Column(EncryptedType( sa.Unicode, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) access_token = sa.Column(EncryptedType( sa.String, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) is_active = sa.Column(EncryptedType( sa.Boolean, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) accounts_num = sa.Column(EncryptedType( sa.Integer, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) phone = sa.Column(EncryptedType( PhoneNumberType, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) color = sa.Column(EncryptedType( ColorType, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) date = sa.Column(EncryptedType( sa.Date, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) time = sa.Column(EncryptedType( sa.Time, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) datetime = sa.Column(EncryptedType( sa.DateTime, test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) enum = sa.Column(EncryptedType( sa.Enum('One', name='user_enum_t'), test_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) return User @@ -217,7 +227,7 @@ def date_simple(): class EncryptedTypeTestCase(object): @pytest.fixture - def Team(self, Base, encryption_engine): + def Team(self, Base, encryption_engine, padding_mechanism): self._team_key = None class Team(Base): @@ -227,7 +237,8 @@ class Team(Base): name = sa.Column(EncryptedType( sa.Unicode, lambda: self._team_key, - encryption_engine) + encryption_engine, + padding_mechanism) ) return Team @@ -289,9 +300,6 @@ def test_lookup_key(self, session, Team): assert team.name == u'One' - with pytest.raises(Exception): - session.query(Team).get(team_2_id) - session.expunge_all() self._team_key = session.query(Team.key).filter_by( @@ -302,9 +310,6 @@ def test_lookup_key(self, session, Team): assert team.name == u'Two' - with pytest.raises(Exception): - session.query(Team).get(team_1_id) - session.expunge_all() # Remove teams @@ -312,7 +317,7 @@ def test_lookup_key(self, session, Team): session.commit() -class TestAesEncryptedTypeTestcase(EncryptedTypeTestCase): +class AesEncryptedTypeTestCase(EncryptedTypeTestCase): @pytest.fixture def encryption_engine(self): @@ -325,6 +330,34 @@ def test_lookup_by_encrypted_string(self, session, User, user, user_name): assert test.username == user.username + +class TestAesEncryptedTypeWithPKCS5Padding(AesEncryptedTypeTestCase): + + @pytest.fixture + def padding_mechanism(self): + return 'pkcs5' + + +class TestAesEncryptedTypeWithOneAndZeroesPadding(AesEncryptedTypeTestCase): + + @pytest.fixture + def padding_mechanism(self): + return 'oneandzeroes' + + +class TestAesEncryptedTypeWithZeroesPadding(AesEncryptedTypeTestCase): + + @pytest.fixture + def padding_mechanism(self): + return 'zeroes' + + +class TestAesEncryptedTypeTestcaseWithNaivePadding(AesEncryptedTypeTestCase): + + @pytest.fixture + def padding_mechanism(self): + return 'naive' + def test_decrypt_raises_value_error_with_invalid_key(self, session, Team): self._team_key = 'one' team = Team(key=self._team_key, name=u'One') @@ -342,6 +375,10 @@ class TestFernetEncryptedTypeTestCase(EncryptedTypeTestCase): def encryption_engine(self): return FernetEngine + @pytest.fixture + def padding_mechanism(self): + return None + class TestDatetimeHandler(object):