desec-stack/api/desecapi/models/mfa.py

from __future__ import annotations

import base64
from functools import cached_property
import secrets
import uuid

from django.conf import settings
from django.db import models, transaction
from django.utils import timezone
from pyotp import TOTP, utils as pyotp_utils


class BaseFactor(models.Model):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    user = models.ForeignKey("User", on_delete=models.CASCADE)
    created = models.DateTimeField(auto_now_add=True)
    last_used = models.DateTimeField(null=True, blank=True)
    name = models.CharField(blank=True, default="", max_length=64)

    class Meta:
        constraints = [
            models.UniqueConstraint(fields=["user", "name"], name="unique_user_name"),
        ]

    @transaction.atomic()
    def delete(self):
        if self.last_used is not None:
            self.user.save(credentials_changed=True)
        return super().delete()

    @transaction.atomic()
    def save(self, *args, **kwargs):
        if not self.user.mfa_enabled:  # enabling MFA
            self.user.save(credentials_changed=True)
        return super().save(*args, **kwargs)


class TOTPFactor(BaseFactor):
    @staticmethod
    def _secret_default():
        return secrets.token_bytes(32)

    secret = models.BinaryField(max_length=32, default=_secret_default.__func__)
    last_verified_timestep = models.PositiveIntegerField(default=0)

    @cached_property
    def _totp(self):
        # TODO switch to self.secret once https://github.com/pyauth/pyotp/pull/138 is released
        return TOTP(self.base32_secret, digits=6)

    @property
    def base32_secret(self):
        return base64.b32encode(self.secret).rstrip(b"=").decode("ascii")

    @property
    def uri(self):
        return self._totp.provisioning_uri(
            name=self.name,
            issuer_name=f"desec.{settings.DESECSTACK_DOMAIN}",
        )

    @transaction.atomic
    def verify(self, code):
        now = timezone.now()
        timestep_now = self._totp.timecode(now)

        for offset in (-1, 0, 1):
            timestep = timestep_now + offset
            if not (self.last_verified_timestep < timestep):
                continue
            if pyotp_utils.strings_equal(str(code), self._totp.generate_otp(timestep)):
                self.last_used = now
                self.last_verified_timestep = timestep
                self.save()
                return True
        return False