You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
57 lines
1.8 KiB
57 lines
1.8 KiB
#!/usr/bin/env python
|
|
# -*- coding:utf-8 -*-
|
|
# project: 11月
|
|
# author: NinEveN
|
|
# date: 2021/11/3
|
|
from django.conf import settings
|
|
from django.db import models
|
|
|
|
from api.utils.baseutils import AESCipher
|
|
|
|
|
|
class AESCharField(models.CharField):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
if 'prefix' in kwargs:
|
|
self.prefix = kwargs['prefix']
|
|
del kwargs['prefix']
|
|
else:
|
|
self.prefix = "aes:::"
|
|
self.cipher = AESCipher(settings.SECRET_KEY)
|
|
super(AESCharField, self).__init__(*args, **kwargs)
|
|
|
|
def deconstruct(self):
|
|
name, path, args, kwargs = super(AESCharField, self).deconstruct()
|
|
if self.prefix != "aes:::":
|
|
kwargs['prefix'] = self.prefix
|
|
return name, path, args, kwargs
|
|
|
|
def from_db_value(self, value, *args, **kwargs):
|
|
if value is None:
|
|
return value
|
|
if value.startswith(self.prefix):
|
|
value = value[len(self.prefix):]
|
|
if isinstance(value, str):
|
|
value = value.encode('utf-8')
|
|
value = self.cipher.decrypt(value)
|
|
return value
|
|
|
|
def to_python(self, value):
|
|
if value is None:
|
|
return value
|
|
elif value.startswith(self.prefix):
|
|
value = value[len(self.prefix):]
|
|
if isinstance(value, str):
|
|
value = value.encode('utf-8')
|
|
value = self.cipher.decrypt(value)
|
|
return value
|
|
|
|
def get_prep_value(self, value):
|
|
if isinstance(value, str):
|
|
value = value.encode('utf-8')
|
|
if isinstance(value, bytes):
|
|
value = self.cipher.encrypt(value)
|
|
value = self.prefix + value.decode('utf-8')
|
|
elif value is not None:
|
|
raise TypeError(str(value) + " is not a valid value for AESCharField")
|
|
return value
|
|
|