Merge pull request #301 from konstantinoskostis/fix/encryption_decryp…

…tion_padding

Fix/encryption decryption padding

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
-

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

sqlalchemy_utils/types/encrypted/__init__.py

@@ -0,0 +1 @@
+# Module for encrypted type

sqlalchemy_utils/types/encrypted.py → 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,20 +67,29 @@ 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):
             value = repr(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,33 +144,66 @@ 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
         engine = create_engine('sqlite:///:memory:')
         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):

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
+}