diff --git a/apps/common/fields.py b/apps/common/fields.py index a06106cfa..355344cd9 100644 --- a/apps/common/fields.py +++ b/apps/common/fields.py @@ -2,11 +2,15 @@ # import json +from django.db import models from django import forms from django.utils import six from django.core.exceptions import ValidationError from django.utils.translation import ugettext as _ from rest_framework import serializers +from .utils import get_signer + +signer = get_signer() class DictField(forms.Field): @@ -46,4 +50,27 @@ class StringIDField(serializers.Field): class StringManyToManyField(serializers.RelatedField): def to_representation(self, value): - return value.__str__() \ No newline at end of file + return value.__str__() + + +class EncryptMixin: + def from_db_value(self, value, expression, connection, context): + if value is not None: + return signer.unsign(value) + return super().from_db_value(self, value, expression, connection, context) + + def get_prep_value(self, value): + if value is None: + return value + return signer.sign(value).decode('utf-8') + + +class EncryptTextField(EncryptMixin, models.TextField): + description = _("Encrypt field using Secret Key") + + +class EncryptCharField(EncryptMixin, models.CharField): + def __init__(self, *args, **kwargs): + kwargs['max_length'] = 2048 + super().__init__(*args, **kwargs) + diff --git a/apps/users/models/user.py b/apps/users/models/user.py index cb8f72674..6c3b467e9 100644 --- a/apps/users/models/user.py +++ b/apps/users/models/user.py @@ -45,7 +45,7 @@ class User(AbstractUser): wechat = models.CharField(max_length=128, blank=True, verbose_name=_('Wechat')) phone = models.CharField(max_length=20, blank=True, null=True, verbose_name=_('Phone')) otp_level = models.SmallIntegerField(default=0, choices=OTP_LEVEL_CHOICES, verbose_name=_('Enable OTP')) - otp_secret_key = models.CharField(max_length=16, blank=True, null=True) + _otp_secret_key = models.CharField(max_length=128, blank=True, null=True) # Todo: Auto generate key, let user download _private_key = models.CharField(max_length=5000, blank=True, verbose_name=_('Private key')) _public_key = models.CharField(max_length=5000, blank=True, verbose_name=_('Public key')) @@ -70,6 +70,14 @@ class User(AbstractUser): def password_raw(self, password_raw_): self.set_password(password_raw_) + @property + def otp_secret_key(self): + return signer.unsign(self._otp_secret_key) + + @otp_secret_key.setter + def otp_secret_key(self, item): + self._otp_secret_key = signer.sign(item).decode('utf-8') + def get_absolute_url(self): return reverse('users:user-detail', args=(self.id,))