SCloud_Django/account/models.py

from datetime import datetime
from django.core.mail import send_mail
from django.db import models
from django.utils.translation import gettext_lazy as _
from django.utils.http import base36_to_int, int_to_base36
from django.conf import settings
from django.utils.crypto import salted_hmac

from .utils import encode_password
from .validators import ASCIIUsernameValidator


class User(models.Model):
    username = models.CharField(
        max_length=25,
        unique=True,
        validators=[ASCIIUsernameValidator()]
    )
    password = models.CharField(_('password'), max_length=128)
    last_login = models.DateTimeField(_('last login'), blank=True, null=True)
    email = models.EmailField(_('email address'), unique=True)

    def save(self, *args, **kwargs):
        self.password = encode_password(self.username, self.password)
        super(User, self).save(*args, **kwargs)

    def get_root_folder(self):
        return self.folders.get(father_folder=None, group=None)

    def send_email(self, subject, message, from_email=None, **kwargs):
        send_mail(subject, message, from_email, [self.email], **kwargs)

    def make_token(self):
        return self._make_token(_timestamp())

    def check_token(self, token):
        if not token:
            return False
        try:
            ts_b36, hash_str = token.split('-')
        except ValueError:
            return False

        try:
            ts = base36_to_int(ts_b36)
        except ValueError:
            return False

        if self._make_token(ts) != token:
            return False

        timestamp = _timestamp()
        if (timestamp - ts) > settings.PASSWORD_RESET_TIMEOUT:
            return False

        return True

    def _make_token(self, timestamp):
        ts_b36 = int_to_base36(timestamp)
        salt = settings.SALT
        value = self._make_hash_value(timestamp)
        secret = settings.SECRET_KEY
        hash_str = salted_hmac(
            salt, value, secret=secret, algorithm='sha256'
        ).hexdigest()[::2]
        token = "%s-%s" % (ts_b36, hash_str)
        return token

    def _make_hash_value(self, timestamp):
        return f'{self.pk}{self.password}{timestamp}{self.email}'

    def __str__(self):
        return self.username


def get_user(request):
    if hasattr(request, 'user') and isinstance(request.user, User):
        print(f'get user from request.user, username={request.user.username}')
        return request.user
    username = request.POST.get('username', '')
    token = request.POST.get('token', '')
    try:
        user = User.objects.get(username=username)
        if user.check_token(token):
            user.tokens.get(token=token)
            return user
    except:
        return None
    return None


def _timestamp():
    dt = datetime.now()
    return int((dt - datetime(2001, 1, 1)).total_seconds())


class LoginToken(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE, related_name='tokens')
    token = models.CharField(max_length=256)