Adunatio/internal_lib/EncryptedField.py

import binascii
import re

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):
        import os
        priv_key = os.environ.get('ADUNATIO_PRIV_KEY')
        try:
            self.keyPair = RSA.importKey(open(priv_key).read())
        except:
            pass

        super().__init__(**kwargs)

    def __get__(self, instance, owner):
        import binascii
        if instance:
            value = instance._data.get(self.name)
            if value:
                encryptor = PKCS1_OAEP.new(self.keyPair)
                return encryptor.decrypt(binascii.unhexlify(value)).decode("utf-8")
            else:
                return None

    def __set__(self, instance, value):
        super().__set__(instance, value)
        if value is not None:
            key = self.name
            try:
                encryptor = PKCS1_OAEP.new(self.keyPair.publickey())
                instance._data[key] = binascii.hexlify(encryptor.encrypt(value.encode('utf-8')))
                instance._mark_as_changed(key)
            except:
                instance._data[key] = 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 prepare_query_value(self, op, value):
        return super().prepare_query_value(op, value)