import base64 import string import json from random import SystemRandom import bcrypt import resumablehashlib from peewee import TextField, CharField, SmallIntegerField from data.text import prefix_search def random_string(length=16): random = SystemRandom() return ''.join([random.choice(string.ascii_uppercase + string.digits) for _ in range(length)]) class _ResumableSHAField(TextField): def _create_sha(self): raise NotImplementedError def db_value(self, value): if value is None: return None sha_state = value.state() # One of the fields is a byte string, let's base64 encode it to make sure # we can store and fetch it regardless of default collocation. sha_state[3] = base64.b64encode(sha_state[3]) return json.dumps(sha_state) def python_value(self, value): if value is None: return None sha_state = json.loads(value) # We need to base64 decode the data bytestring. sha_state[3] = base64.b64decode(sha_state[3]) to_resume = self._create_sha() to_resume.set_state(sha_state) return to_resume class ResumableSHA256Field(_ResumableSHAField): def _create_sha(self): return resumablehashlib.sha256() class ResumableSHA1Field(_ResumableSHAField): def _create_sha(self): return resumablehashlib.sha1() class JSONField(TextField): def db_value(self, value): return json.dumps(value) def python_value(self, value): if value is None or value == "": return {} return json.loads(value) class Base64BinaryField(TextField): def db_value(self, value): if value is None: return None return base64.b64encode(value) def python_value(self, value): if value is None: return None return base64.b64decode(value) class DecryptedValue(object): """ Wrapper around an already decrypted value to be placed into an encrypted field. """ def __init__(self, decrypted_value): assert decrypted_value is not None self.value = decrypted_value def decrypt(self): return self.value def matches(self, unencrypted_value): """ Returns whether the value of this field matches the unencrypted_value. """ return self.decrypt() == unencrypted_value class LazyEncryptedValue(object): """ Wrapper around an encrypted value in an encrypted field. Will decrypt lazily. """ def __init__(self, encrypted_value, field): self.encrypted_value = encrypted_value self._field = field def decrypt(self): """ Decrypts the value. """ return self._field.model._meta.encrypter.decrypt_value(self.encrypted_value) def matches(self, unencrypted_value): """ Returns whether the value of this field matches the unencrypted_value. """ return self.decrypt() == unencrypted_value def __eq__(self, _): raise Exception('Disallowed operation; use `matches`') def __mod__(self, _): raise Exception('Disallowed operation; use `matches`') def __pow__(self, _): raise Exception('Disallowed operation; use `matches`') def __contains__(self, _): raise Exception('Disallowed operation; use `matches`') def contains(self, _): raise Exception('Disallowed operation; use `matches`') def startswith(self, _): raise Exception('Disallowed operation; use `matches`') def endswith(self, _): raise Exception('Disallowed operation; use `matches`') def _add_encryption(field_class, requires_length_check=True): """ Adds support for encryption and decryption to the given field class. """ class indexed_class(field_class): def __init__(self, default_token_length=None, *args, **kwargs): def _generate_default(): return DecryptedValue(random_string(default_token_length)) if default_token_length is not None: kwargs['default'] = _generate_default field_class.__init__(self, *args, **kwargs) assert not self.index def db_value(self, value): if value is None: return None if isinstance(value, LazyEncryptedValue): return value.encrypted_value if isinstance(value, DecryptedValue): value = value.value meta = self.model._meta return meta.encrypter.encrypt_value(value, self.max_length if requires_length_check else None) def python_value(self, value): if value is None: return None return LazyEncryptedValue(value, self) def __eq__(self, _): raise Exception('Disallowed operation; use `matches`') def __mod__(self, _): raise Exception('Disallowed operation; use `matches`') def __pow__(self, _): raise Exception('Disallowed operation; use `matches`') def __contains__(self, _): raise Exception('Disallowed operation; use `matches`') def contains(self, _): raise Exception('Disallowed operation; use `matches`') def startswith(self, _): raise Exception('Disallowed operation; use `matches`') def endswith(self, _): raise Exception('Disallowed operation; use `matches`') return indexed_class EncryptedCharField = _add_encryption(CharField) EncryptedTextField = _add_encryption(TextField, requires_length_check=False) class EnumField(SmallIntegerField): def __init__(self, enum_type, *args, **kwargs): kwargs.pop('index', None) super(EnumField, self).__init__(index=True, *args, **kwargs) self.enum_type = enum_type def db_value(self, value): """Convert the python value for storage in the database.""" return int(value.value) def python_value(self, value): """Convert the database value to a pythonic value.""" return self.enum_type(value) if value is not None else None def clone_base(self, **kwargs): return super(EnumField, self).clone_base( enum_type=self.enum_type, **kwargs) def _add_fulltext(field_class): """ Adds support for full text indexing and lookup to the given field class. """ class indexed_class(field_class): # Marker used by SQLAlchemy translation layer to add the proper index for full text searching. __fulltext__ = True def __init__(self, match_function, *args, **kwargs): field_class.__init__(self, *args, **kwargs) self.match_function = match_function def match(self, query): return self.match_function(self, query) def match_prefix(self, query): return prefix_search(self, query) def __mod__(self, _): raise Exception('Unsafe operation: Use `match` or `match_prefix`') def __pow__(self, _): raise Exception('Unsafe operation: Use `match` or `match_prefix`') def __contains__(self, _): raise Exception('Unsafe operation: Use `match` or `match_prefix`') def contains(self, _): raise Exception('Unsafe operation: Use `match` or `match_prefix`') def startswith(self, _): raise Exception('Unsafe operation: Use `match` or `match_prefix`') def endswith(self, _): raise Exception('Unsafe operation: Use `match` or `match_prefix`') return indexed_class FullIndexedCharField = _add_fulltext(CharField) FullIndexedTextField = _add_fulltext(TextField) class Credential(object): """ Credential represents a hashed credential. """ def __init__(self, hashed): self.hashed = hashed def matches(self, value): """ Returns true if this credential matches the unhashed value given. """ return bcrypt.hashpw(value.encode('utf-8'), self.hashed) == self.hashed @classmethod def from_string(cls, string_value): """ Returns a Credential object from an unhashed string value. """ return Credential(bcrypt.hashpw(string_value.encode('utf-8'), bcrypt.gensalt())) @classmethod def generate(cls, length=20): """ Generates a new credential and returns it, along with its unhashed form. """ token = random_string(length) return Credential.from_string(token), token class CredentialField(CharField): """ A character field that stores crytographically hashed credentials that should never be available to the user in plaintext after initial creation. This field automatically provides verification. """ def __init__(self, *args, **kwargs): CharField.__init__(self, *args, **kwargs) assert 'default' not in kwargs assert not self.index def db_value(self, value): if value is None: return None if isinstance(value, basestring): raise Exception('A string cannot be given to a CredentialField; please wrap in a Credential') return value.hashed def python_value(self, value): if value is None: return None return Credential(value)