diff --git a/djblets/db/fields/base64_field.py b/djblets/db/fields/base64_field.py index 1c49d3fa6210d430cf5a9942a4eecca08afdb073..e2e3fd7da86ed2e2ba40e1d30b2908f13e08bb3f 100644 --- a/djblets/db/fields/base64_field.py +++ b/djblets/db/fields/base64_field.py @@ -1,81 +1,262 @@ +"""Field for storing data as Base64-encoded values. + +See :py:class:`Base64Field` for the main field type and usage instructions. +""" + from __future__ import unicode_literals -import base64 +try: + # Python >= 3.1 + from base64 import (decodebytes as base64_decode, + encodebytes as base64_encode) +except ImportError: + # Python < 3.1 + from base64 import (decodestring as base64_decode, + encodestring as base64_encode) from django.db import models +from django.utils import six from django.utils.encoding import smart_text +from django.utils.translation import ugettext as _ -class Base64DecodedValue(str): - """ - A subclass of string that can be identified by Base64Field, in order - to prevent double-encoding or double-decoding. +class Base64TypeError(TypeError): + """Error indicating an invalid value type was provided for the field.""" + + def __init__(self, value): + """Initialize the error. + + Args: + value (object): + The invalid value that was passed. + """ + super(Base64TypeError, self).__init__( + _('The provided value must be a string type or a ' + 'Base64DecodedValue, not %r.') + % type(value)) + + +class Base64DecodedValue(bytes): + """An identifiable wrapper around byte string values for Base64Field. + + This wraps any values coming from :py:class:`Base64Field`, helping make + a distinction between arbitrary strings and decoded strings from the + database. It helps prevent double-encoding or double-decoding of data. """ - pass class Base64FieldCreator(object): + """Property-like class used to store/retrieve Base64 values. + + This works much like a property, and takes care of encoding strings for + storage and decoding them on retrieval. It's set internally by + :py:class:`Base64Field` in place of the normal field attribute. + """ + def __init__(self, field): + """Initialize the creator. + + Args: + field (Base64Field): + The field owning this object. + """ self.field = field def __set__(self, obj, value): - pk_val = obj._get_pk_val(obj.__class__._meta) - pk_set = pk_val is not None and smart_text(pk_val) != '' + """Set a new value in the field. + + If this is a :py:class:`Base64DecodedValue`, or the model is new and + hasn't yet been persisted to the database, the value will be encoded + and stored. Otherwise, if it's a standard string value or the model + is not new, it's assumed that this is encoded data for storage, and + will be stored directly. + + Args: + obj (django.db.models.Model): + The model owning the field. - if (isinstance(value, Base64DecodedValue) or not pk_set): - obj.__dict__[self.field.name] = base64.encodestring(value) - else: - obj.__dict__[self.field.name] = value + value (object): + The value being set. This must be a valid string value or + :py:class:`Base64DecodedValue`. - setattr(obj, "%s_initted" % self.field.name, True) + Raises: + Base64TypeError: + The type of value provided could not be set. + """ + if value is not None: + if not isinstance(value, (six.text_type, six.binary_type)): + raise Base64TypeError(value) - def __get__(self, obj, type=None): + pk_val = obj._get_pk_val(obj.__class__._meta) + pk_set = pk_val is not None and smart_text(pk_val) != '' + + if isinstance(value, six.text_type): + value = value.encode('utf-8') + + if isinstance(value, Base64DecodedValue) or not pk_set: + value = base64_encode(value) + + obj.__dict__[self.field.name] = value + + def __get__(self, obj, *args, **kwargs): + """Return a decoded value from the field. + + Args: + obj (django.db.models.Model): + The model owning the field. + + *args (tuple): + Unused positional arguments. + + **kwargs (dict): + Unused keyword arguments. + + Returns: + Base64DecodedValue: + The decoded value from the field. If no value has yet been stored, + this will return ``None`` instead. + + Raises: + AttributeError: + A ``None`` value was passed for the object. + """ if obj is None: raise AttributeError('Can only be accessed via an instance.') value = obj.__dict__[self.field.name] - if value is None: - return None - else: - return Base64DecodedValue(base64.decodestring(value)) + if value is not None: + value = Base64DecodedValue(base64_decode(value)) + + return value class Base64Field(models.TextField): + """A text field for storing Base64-encoded values. + + This is used to store data (such as binary data or encoding-sensitive data) + to the database in a Base64 encoding. This is useful if you're dealing + with unknown encodings and must guarantee that no modifications to the text + occurs and that you can read/write the data in any database with any + encoding. + + When accessing this field on an instance of a model, a + :py:class:`Base64DecodedValue` will be returned consisting of the decoded + data. This is a byte string, and can be treated as such. If set back into + the field, it will be re-encoded and stored. + + When writing to the field, the behavior changes based on the type of value + and the state of the model: + + * If the model instance is new (has not yet been saved in the database), + any string set will be encoded. This allows the value to be passed during + a ``create()`` call. + + * If the model is not new, any string that's set will be assumed to be + encoded by the caller. + + * Passing a :py:class:`Base64DecodedValue` byte string will always cause + the stored data to be encoded. + + The field also adds a :samp:`get_{fieldname}_base64()` method to the + class, which returns the raw Base64 encoded content from the database. """ - A subclass of TextField that encodes its data as base64 in the database. - This is useful if you're dealing with unknown encodings and must guarantee - that no modifications to the text occurs and that you can read/write - the data in any database with any encoding. - """ - serialize_to_string = True def contribute_to_class(self, cls, name): + """Set attributes on a new model class. + + This is called when constructing a model class making use of this + field. It sets the field's attribute to a + :py:class:`Base64FieldCreator` and adds a + :samp:`get_{fieldname}_base64()` method to the class. + + Args: + cls (type): + The class to add the arguments to. + + name (str): + The name of the field. + """ super(Base64Field, self).contribute_to_class(cls, name) setattr(cls, self.name, Base64FieldCreator(self)) setattr(cls, 'get_%s_base64' % self.name, lambda model_instance: model_instance.__dict__[self.name]) - def get_db_prep_value(self, value, connection=None, prepared=False): - if isinstance(value, Base64DecodedValue): - value = base64.encodestring(value) + def get_prep_value(self, value): + """Return a value prepared for the field. - return value + Args: + value (object): + The value to prepare. This is expected to be a string or a + :py:class:`Base64DecodedValue`. If the latter, it will be + encoded. + + Returns: + bytes: + The resulting value. - def save_form_data(self, instance, data): - setattr(instance, self.name, Base64DecodedValue(data)) + Raises: + Base64TypeError: + The type of value provided could not be prepared for writing. + """ + if value is not None: + if not isinstance(value, (six.text_type, six.binary_type)): + raise Base64TypeError(value) + + if isinstance(value, Base64DecodedValue): + value = base64_encode(value) + elif isinstance(value, six.text_type): + value = value.encode('utf-8') + + return value def to_python(self, value): - if isinstance(value, Base64DecodedValue): - return value - else: - return Base64DecodedValue(base64.decodestring(value)) + """Return a Python representation of a value for the field. + + This will decode the value (if not already decoded) and return it. + + Args: + value (object): + The value to return a decoded value for. + + Returns: + Base64DecodedValue: + The decoded version of the provided value. + + Raises: + Base64TypeError: + The type of value provided could not be prepared for writing. + """ + if value is not None: + if not isinstance(value, (six.text_type, six.binary_type)): + raise Base64TypeError(value) + + if not isinstance(value, Base64DecodedValue): + if isinstance(value, six.text_type): + value = value.encode('utf-8') + + value = Base64DecodedValue(base64_decode(value)) + + return value def value_to_string(self, obj): - value = self._get_val_from_obj(obj) + """Return a string representation of the value from a model. + + The returned value will be a Base64-encoded string value. + + Args: + obj (django.db.models.Model): + The model instance owning the field and value. - if isinstance(value, Base64DecodedValue): - return base64.encodestring(value) - else: - return value + Returns: + bytes: + The Base64-encoded byte string for the stored value. + """ + value = self.value_from_object(obj) + assert value is None or isinstance(value, Base64DecodedValue) + + if value is not None: + value = base64_encode(value) + + return value diff --git a/djblets/db/tests/test_base64_field.py b/djblets/db/tests/test_base64_field.py new file mode 100644 index 0000000000000000000000000000000000000000..ad40c417fe670fc7960741a92f1ef70d8fb907a0 --- /dev/null +++ b/djblets/db/tests/test_base64_field.py @@ -0,0 +1,284 @@ +# coding: utf-8 +"""Unit tests for djblets.db.fields.base64_field.""" + +from __future__ import unicode_literals + +from django.db import models + +from djblets.db.fields.base64_field import (Base64DecodedValue, Base64Field, + Base64TypeError) +from djblets.testing.testcases import TestCase, TestModelsLoaderMixin + + +class Base64TestModel(models.Model): + field = Base64Field(null=True) + + +class Base64FieldTests(TestModelsLoaderMixin, TestCase): + """Unit tests for djblets.db.fields.base64_field.Base64Field.""" + + tests_app = 'djblets.db.tests' + + def test_create_with_decoded_bytes_value(self): + """Testing Base64Field with setting decoded bytes value in + Model.objects.create() + """ + obj = Base64TestModel.objects.create(field=b'This is a test') + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_create_with_decoded_unicode_value(self): + """Testing Base64Field with setting decoded unicode value in + Model.objects.create() + """ + obj = Base64TestModel.objects.create(field='This is a tést') + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a t\xc3\xa9st') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHTDqXN0\n') + + def test_create_with_base64_decoded_value(self): + """Testing Base64Field with setting Base64DecodedValue in + Model.objects.create() + """ + obj = Base64TestModel.objects.create( + field=Base64DecodedValue(b'This is a test')) + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_create_with_none(self): + """Testing Base64Field with setting None in + Model.objects.create() + """ + obj = Base64TestModel.objects.create(field=None) + self.assertIsNone(obj.field) + + encoded = obj.get_field_base64() + self.assertIsNone(encoded) + + def test_create_with_invalid_value_type(self): + """Testing Base64Field with setting invalid value type in + Model.objects.create() + """ + with self.assertRaises(Base64TypeError): + Base64TestModel.objects.create(field=True) + + def test_unsaved_obj_with_decoded_bytes_value(self): + """Testing Base64Field with setting decoded bytes value on unsaved + instance + """ + obj = Base64TestModel() + obj.field = b'This is a test' + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_unsaved_obj_with_decoded_unicode_value(self): + """Testing Base64Field with setting decoded unicode value on unsaved + instance + """ + obj = Base64TestModel() + obj.field = 'This is a tést' + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a t\xc3\xa9st') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHTDqXN0\n') + + def test_unsaved_obj_with_base64_decoded_value(self): + """Testing Base64Field with setting Base64DecodedValue on unsaved + instance + """ + obj = Base64TestModel() + obj.field = Base64DecodedValue(b'This is a test') + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_unsaved_obj_with_invalid_value_type(self): + """Testing Base64Field with setting invalid value type on unsaved + instance + """ + obj = Base64TestModel() + + with self.assertRaises(Base64TypeError): + obj.field = True + + def test_unsaved_obj_with_none(self): + """Testing Base64Field with setting None on unsaved instance""" + obj = Base64TestModel() + obj.field = None + + self.assertIsNone(obj.field) + + encoded = obj.get_field_base64() + self.assertIsNone(encoded) + + def test_saved_obj_with_decoded_bytes_value(self): + """Testing Base64Field with setting decoded bytes value on saved + instance + """ + obj = Base64TestModel.objects.create() + obj.field = b'VGhpcyBpcyBhIHRlc3Q=\n' + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_saved_obj_with_decoded_unicode_value(self): + """Testing Base64Field with setting decoded unicode value on saved + instance + """ + obj = Base64TestModel.objects.create() + obj.field = 'VGhpcyBpcyBhIHTDqXN0\n' + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a t\xc3\xa9st') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHTDqXN0\n') + + def test_saved_obj_with_base64_decoded_value(self): + """Testing Base64Field with setting Base64DecodedValue on saved + instance + """ + obj = Base64TestModel.objects.create() + obj.field = Base64DecodedValue(b'This is a test') + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_saved_obj_with_invalid_value_type(self): + """Testing Base64Field with setting invalid value type on saved + instance + """ + obj = Base64TestModel.objects.create() + + with self.assertRaises(Base64TypeError): + obj.field = True + + def test_saved_obj_with_none(self): + """Testing Base64Field with setting None on saved instance""" + obj = Base64TestModel.objects.create() + obj.field = None + + self.assertIsNone(obj.field) + + encoded = obj.get_field_base64() + self.assertIsNone(encoded) + + def test_values_persist_after_save(self): + """Testing Base64Field value persists correctly after save""" + obj = Base64TestModel.objects.create(field=b'This is a test') + obj = Base64TestModel.objects.get(pk=obj.pk) + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_save_form_data_with_bytes(self): + """Testing Base64Field.save_form_data with bytes value""" + obj = Base64TestModel(field=b'This is a test') + obj._meta.get_field('field').save_form_data(obj, b'This is a test') + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a test') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_save_form_data_with_unicode(self): + """Testing Base64Field.save_form_data with unicode value""" + obj = Base64TestModel(field=b'This is a test') + obj._meta.get_field('field').save_form_data(obj, 'This is a tést') + + self.assertIs(type(obj.field), Base64DecodedValue) + self.assertEqual(obj.field, b'This is a t\xc3\xa9st') + + encoded = obj.get_field_base64() + self.assertIs(type(encoded), bytes) + self.assertEqual(encoded, b'VGhpcyBpcyBhIHTDqXN0\n') + + def test_to_python_with_bytes(self): + """Testing Base64Field.to_python with bytes value""" + obj = Base64TestModel() + value = obj._meta.get_field('field').to_python( + b'VGhpcyBpcyBhIHRlc3Q=\n') + + self.assertIs(type(value), Base64DecodedValue) + self.assertEqual(value, b'This is a test') + + def test_to_python_with_unicode(self): + """Testing Base64Field.to_python with unicode value""" + obj = Base64TestModel() + value = obj._meta.get_field('field').to_python( + 'VGhpcyBpcyBhIHRlc3Q=\n') + + self.assertIs(type(value), Base64DecodedValue) + self.assertEqual(value, b'This is a test') + + def test_to_python_with_base64_decoded_value(self): + """Testing Base64Field.to_python with Base64DecodedValue""" + obj = Base64TestModel() + value = obj._meta.get_field('field').to_python( + Base64DecodedValue(b'This is a test')) + + self.assertIs(type(value), Base64DecodedValue) + self.assertEqual(value, b'This is a test') + + def test_to_python_with_none(self): + """Testing Base64Field.to_python with None""" + obj = Base64TestModel() + value = obj._meta.get_field('field').to_python(None) + + self.assertIsNone(value) + + def test_value_to_string_with_base64_decoded_value(self): + """Testing Base64Field.value_to_string with Base64DecodedValue""" + obj = Base64TestModel(field=b'This is a test') + value = obj._meta.get_field('field').value_to_string(obj) + + self.assertIs(type(value), bytes) + self.assertEqual(value, b'VGhpcyBpcyBhIHRlc3Q=\n') + + def test_value_to_string_with_none(self): + """Testing Base64Field.value_to_string with None""" + obj = Base64TestModel(field=None) + value = obj._meta.get_field('field').value_to_string(obj) + + self.assertIsNone(value)