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')
        self.keyPair = RSA.importKey(open(priv_key).read())

        super().__init__(**kwargs)

    def __get__(self, instance, owner):
        from flask import current_app as app
        app.logger.error(self.name)
        import  binascii
        if instance:
            value = instance._data.get(self.name)
            if value:
                app.logger.error(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:
            from flask import current_app as app
            app.logger.error(self.name)
            app.logger.error(instance._data)
            app.logger.error("data : {} ".format(value))
            print(self.name)
            print(instance._data)
            print(value)
            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 lookup_member(self, member_name):
        return None

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