diff --git a/internal_lib/EncryptedField.py b/internal_lib/EncryptedField.py index 5adf1c4..5d15061 100644 --- a/internal_lib/EncryptedField.py +++ b/internal_lib/EncryptedField.py @@ -1,26 +1,35 @@ import re +from flask import current_app from mongoengine.base import BaseField -from mongoengine.queryset import STRING_OPERATORS +from Crypto.PublicKey import RSA +from Crypto.Cipher import PKCS1_OAEP class EncryptedStringField(BaseField): - """A unicode string field.""" + """A unicode encrypted string field.""" + keyPair = None + + def __init__(self, **kwargs): + + self.keyPair = RSA.importKey(open("privkey.pem").read()) - def __init__(self, regex=None, max_length=None, min_length=None, **kwargs): - """ - :param regex: (optional) A string pattern that will be applied during validation - :param max_length: (optional) A max length that will be applied during validation - :param min_length: (optional) A min length that will be applied during validation - :param kwargs: Keyword arguments passed into the parent :class:`~mongoengine.BaseField` - """ - self.regex = re.compile(regex) if regex else None - self.max_length = max_length - self.min_length = min_length super().__init__(**kwargs) - def to_python(self, value): + def __get__(self, instance, owner): + value = instance._data.get(self.name) + encryptor = PKCS1_OAEP.new(self.keyPair.publickey()) + return encryptor.decrypt(value) + + def __set__(self, instance, value): + key = self.name + encryptor = PKCS1_OAEP.new(self.keyPair.publickey()) + + instance._data[key] = encryptor.encrypt(value) + instance._mark_as_changed(key) + + def to_python(self, value): if isinstance(value, str): return value try: @@ -29,41 +38,8 @@ class EncryptedStringField(BaseField): pass return value - def validate(self, value): - if not isinstance(value, str): - self.error("StringField only accepts string values") - - if self.max_length is not None and len(value) > self.max_length: - self.error("String value is too long") - - if self.min_length is not None and len(value) < self.min_length: - self.error("String value is too short") - - if self.regex is not None and self.regex.match(value) is None: - self.error("String value did not match validation regex") - def lookup_member(self, member_name): return None def prepare_query_value(self, op, value): - if not isinstance(op, str): - return value - - if op in STRING_OPERATORS: - case_insensitive = op.startswith("i") - op = op.lstrip("i") - - flags = re.IGNORECASE if case_insensitive else 0 - - regex = r"%s" - if op == "startswith": - regex = r"^%s" - elif op == "endswith": - regex = r"%s$" - elif op == "exact": - regex = r"^%s$" - - # escape unsafe characters which could lead to a re.error - value = re.escape(value) - value = re.compile(regex % value, flags) - return super().prepare_query_value(op, value) \ No newline at end of file + return super().prepare_query_value(op, value)