Adunatio/internal_lib/EncryptedField.py

import re

from flask import current_app
from mongoengine.base import BaseField
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP


class EncryptedStringField(BaseField):
    """A unicode encrypted string field."""
    keyPair = None

    def __init__(self, **kwargs):

        self.keyPair = RSA.importKey(open(current_app.config.get('priv_key_path')).read())

        super().__init__(**kwargs)

    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:
            value = value.decode("utf-8")
        except Exception:
            pass
        return value

    def lookup_member(self, member_name):
        return None

    def prepare_query_value(self, op, value):
        return super().prepare_query_value(op, value)