/src/niwi/contrib/db/fields/encrypted.py

https://github.com/jespino/niwi-web · Python · 58 lines · 44 code · 13 blank · 1 comment · 7 complexity · d9295319a48c99ffb669f1fa9eecf025 MD5 · raw file

  1. # -*- coding: utf-8 -*-
  2. from django.db import models
  3. from django.core.exceptions import ImproperlyConfigured
  4. from django import forms
  5. from django.conf import settings
  6. import base64
  7. try:
  8. from Crypto.Cipher import Blowfish
  9. except ImportError:
  10. raise ImportError('Using an encrypted field requires the pycripto. You can install pycripto with "pip install pycrypto"')
  11. class BaseEncryptedField(models.Field):
  12. __prefix = 'string_enc::'
  13. def __init__(self, *args, **kwargs):
  14. self.cipher = Blowfish.new(settings.SECRET_KEY[:10])
  15. super(BaseEncryptedField, self).__init__(*args, **kwargs)
  16. def to_python(self, value):
  17. if value.startswith(self.__prefix):
  18. decrypted_value = self.cipher.decrypt(base64.b64decode(value[len(self.__prefix):]))
  19. rest = decrypted_value.index('\0')
  20. if rest > 0:
  21. decrypted_value = decrypted_value[:rest]
  22. return decrypted_value.decode('utf-8')
  23. else:
  24. return value
  25. def get_db_prep_value(self, value, connection, prepared=False):
  26. if not value.startswith(self.__prefix):
  27. value = value.encode('utf-8')
  28. if len(value) % 8 != 0:
  29. counter = 8 - (len(value) % 8)
  30. value = value + "\0" * counter
  31. return self.__prefix + base64.b64encode(self.cipher.encrypt(value))
  32. return value
  33. class EncryptedTextField(BaseEncryptedField):
  34. __metaclass__ = models.SubfieldBase
  35. def get_internal_type(self):
  36. return 'TextField'
  37. def formfield(self, **kwargs):
  38. defaults = {'widget': forms.Textarea}
  39. defaults.update(kwargs)
  40. return super(EncryptedTextField, self).formfield(**defaults)
  41. def south_field_triple(self):
  42. from south.modelsinspector import introspector
  43. field_class = "django.db.models.fields.TextField"
  44. args, kwargs = introspector(self)
  45. return (field_class, args, kwargs)