Merge pull request #44 from peterfarrell/feature_custom_manager

Feature custom manager

Author: Kévin Etienne

README.md

@@ -127,12 +127,16 @@ INSTALLED_APPS = (
 
 ### Usage
 
-#### Model definition
+#### Model / Manager definition
 
 ```python
 from django.db import models
 
-from pgcrypto import fields
+from pgcrypto import fields, managers
+
+class MyModelManager(managers.PGPManager):
+    pass
+
 
 class MyModel(models.Model):
     digest_field = fields.TextDigestField()
@@ -144,6 +148,8 @@ class MyModel(models.Model):
     pgp_sym_field = fields.TextPGPSymmetricKeyField()
     date_pgp_sym_field = fields.DatePGPSymmetricKeyField()
     datetime_pgp_sym_field = fields.DateTimePGPSymmetricKeyField()
+    
+    objects = MyModelManager()
 ```
 
 #### Encrypting
@@ -155,7 +161,36 @@ Example:
 >>> MyModel.objects.create(value='Value to be encrypted...')
 ```
 
-#### Decrypting
+#### Decryption using custom model managers
+
+If you use the bundled `PGPManager` with your custom model manager, all encrypted 
+fields will automatically decrypted for you (except for hash fields which are one 
+way).
+
+N.B. The bundled manager does not support decryption of fields from FK joins. For 
+example if the `MyModel` class had a FK to to `AnotherModel` class, no encrypted 
+fields be decrypted in the joined `AnotherModel`.
+
+It is recommended that you use the bundled `PGPAdmin` class if using the custom 
+model manager and the Django Admin. The Django Admin performance suffers when 
+using the bundled custom manager. The `PGPAdmin` disables automatic decryption 
+for all ORM calls for that admin class.
+
+```python
+from django.contrib import admin
+
+from pgcrypto.admin import PGPAdmin
+
+
+class MyModelAdmin(admin.ModelAdmin, PGPAdmin):
+    # Your admin code here
+```
+
+
+#### Decrypting using aggregates
+
+This is useful if you are not using the custom manager or need to decrypt fields 
+coming from joined FK fields.
 
 ##### PGP fields
 

pgcrypto/admin.py

@@ -0,0 +1,5 @@
+class PGPAdmin(object):
+
+    def get_queryset(self, request):
+        """Skip any auto decryption when ORM calls are from the admin."""
+        return self.model.objects.get_queryset(**{'skip_decrypt': True})

pgcrypto/managers.py

@@ -0,0 +1,65 @@
+from django.conf import settings
+from django.db import models
+
+
+class PGPManager(models.Manager):
+    use_for_related_fields = True
+    use_in_migrations = True
+
+    @staticmethod
+    def _get_pgp_symmetric_decrypt_sql(field):
+        """Decrypt sql for symmetric fields using the cast sql if required."""
+        sql = """pgp_sym_decrypt("{0}"."{1}", '{2}')"""
+        if hasattr(field, 'cast_sql'):
+            sql = field.cast_sql % sql
+
+        return sql.format(
+            field.model._meta.db_table,
+            field.name,
+            settings.PGCRYPTO_KEY,
+        )
+
+    @staticmethod
+    def _get_pgp_public_key_decrypt_sql(field):
+        """Decrypt sql for public key fields using the cast sql if required."""
+        sql = """pgp_pub_decrypt("{0}"."{1}", dearmor('{2}'))"""
+
+        if hasattr(field, 'cast_sql'):
+            sql = field.cast_sql % sql
+
+        return sql.format(
+            field.model._meta.db_table,
+            field.name,
+            settings.PRIVATE_PGP_KEY,
+        )
+
+    def get_queryset(self, *args, **kwargs):
+        """Decryption in queryset through meta programming."""
+        # importing here otherwise there's a circular reference issue
+        from pgcrypto.mixins import PGPSymmetricKeyFieldMixin, PGPPublicKeyFieldMixin
+
+        skip_decrypt = kwargs.pop('skip_decrypt', None)
+
+        qs = super().get_queryset(*args, **kwargs)
+
+        # The Django admin skips this process because it's extremely slow
+        if not skip_decrypt:
+            select_sql = {}
+            encrypted_fields = []
+
+            for field in self.model._meta.get_fields():
+                if isinstance(field, PGPSymmetricKeyFieldMixin):
+                    select_sql[field.name] = self._get_pgp_symmetric_decrypt_sql(field)
+                    encrypted_fields.append(field.name)
+                elif isinstance(field, PGPPublicKeyFieldMixin):
+                    select_sql[field.name] = self._get_pgp_public_key_decrypt_sql(field)
+                    encrypted_fields.append(field.name)
+
+            # Django queryset.extra() is used here to add decryption sql to query.
+            qs = qs.defer(
+                *encrypted_fields
+            ).extra(
+                select=select_sql
+            )
+
+        return qs

pgcrypto/mixins.py

@@ -83,7 +83,7 @@ class PGPPublicKeyFieldMixin(PGPMixin):
 
 
 class PGPSymmetricKeyFieldMixin(PGPMixin):
-    """PGP symmetric key encrypted field mixin for postgred."""
+    """PGP symmetric key encrypted field mixin for postgres."""
     aggregate = PGPSymmetricKeyAggregate
 
 

tests/models.py

@@ -1,6 +1,10 @@
 from django.db import models
 
-from pgcrypto import fields
+from pgcrypto import fields, managers
+
+
+class EncryptedModelManager(managers.PGPManager):
+    pass
 
 
 class EncryptedModel(models.Model):
@@ -20,3 +24,11 @@ class EncryptedModel(models.Model):
 
     class Meta:
         app_label = 'tests'
+
+
+class EncryptedModelWithManager(EncryptedModel):
+
+    objects = EncryptedModelManager()
+
+    class Meta:
+        proxy = True

tests/test_fields.py

@@ -4,12 +4,10 @@
 from django.test import TestCase
 from incuna_test_utils.utils import field_names
 
-from pgcrypto import aggregates, proxy
-from pgcrypto import fields
+from pgcrypto import aggregates, fields, proxy
 from .factories import EncryptedModelFactory
 from .forms import EncryptedForm
-from .models import EncryptedModel
-
+from .models import EncryptedModel, EncryptedModelWithManager
 
 KEYED_FIELDS = (fields.TextDigestField, fields.TextHMACField)
 EMAIL_PGP_FIELDS = (fields.EmailPGPPublicKeyField, fields.EmailPGPSymmetricKeyField)
@@ -408,3 +406,44 @@ def test_null(self):
         for field in fields:
             with self.subTest(field=field):
                 self.assertEqual(getattr(instance, field), None)
+
+
+class TestPGPManager(TestCase):
+    """Test `PGPManager` can be integrated in a `Django` model."""
+    model = EncryptedModelWithManager
+
+    def test_auto_decryption(self):
+        """Assert auto decryption via manager."""
+        expected_string = 'bonjour'
+        expected_date = date(2016, 9, 1)
+        expected_datetime = datetime(2016, 9, 1, 0, 0, 0)
+
+        EncryptedModelFactory.create(
+            digest_field=expected_string,
+            hmac_field=expected_string,
+            pgp_pub_field=expected_string,
+            pgp_sym_field=expected_string,
+            date_pgp_sym_field=expected_date,  # Tests cast sql
+            datetime_pgp_sym_field=expected_datetime,  # Tests cast sql
+        )
+
+        instance = self.model.objects.get()
+
+        # Using `__dict__` bypasses "on the fly" decryption that normally occurs
+        # if accessing a field that is not yet decrypted.
+        # If decryption is not working, we get references to <In_Memory> classes
+        self.assertEqual(instance.__dict__['pgp_pub_field'], expected_string)
+        self.assertEqual(instance.__dict__['pgp_sym_field'], expected_string)
+        self.assertEqual(instance.__dict__['date_pgp_sym_field'], expected_date)
+        self.assertEqual(instance.__dict__['datetime_pgp_sym_field'], expected_datetime)
+
+        # Ensure digest / hmac fields are unaffected
+        count = self.model.objects.filter(
+            digest_field__hash_of=expected_string
+        ).count()
+        self.assertEqual(count, 1)
+
+        count = self.model.objects.filter(
+            hmac_field__hash_of=expected_string
+        ).count()
+        self.assertEqual(count, 1)