8
0
Fork 0
mirror of https://gitlab2.federez.net/re2o/re2o synced 2024-12-23 07:23:46 +00:00
re2o/preferences/aes_field.py
2018-01-29 16:47:08 +00:00

59 lines
1.6 KiB
Python

import string
import binascii
from random import choice
from Crypto.Cipher import AES
from django.db import models
from django.conf import settings
EOD = '`%EofD%`' # This should be something that will not occur in strings
def genstring(length=16, chars=string.printable):
return ''.join([choice(chars) for i in range(length)])
def encrypt(key, s):
obj = AES.new(key)
datalength = len(s) + len(EOD)
if datalength < 16:
saltlength = 16 - datalength
else:
saltlength = 16 - datalength % 16
ss = ''.join([s, EOD, genstring(saltlength)])
return obj.encrypt(ss)
def decrypt(key, s):
obj = AES.new(key)
ss = obj.decrypt(s)
print(ss)
return ss.split(bytes(EOD, 'utf-8'))[0]
class AESEncryptedField(models.CharField):
def save_form_data(self, instance, data):
if value is None:
return value
setattr(instance, self.name,
binascii.b2a_base64(encrypt(settings.AES_KEY, data)))
def to_python(self, value):
if value is None:
return None
return decrypt(settings.AES_KEY,
binascii.a2b_base64(value)).decode('utf-8')
def from_db_value(self, value, expression, connection, *args):
if value is None:
return value
return decrypt(settings.AES_KEY,
binascii.a2b_base64(value)).decode('utf-8')
def get_prep_value(self, value):
if value is None:
return value
return binascii.b2a_base64(encrypt(
settings.AES_KEY,
value
))