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

from __future__ import annotations

import secrets
import string
import uuid

from django.conf import settings
from django.db import models
from django.utils import timezone
from django_prometheus.models import ExportModelOperationsMixin

from desecapi import metrics


def captcha_default_content(kind: str) -> str:
    if kind == Captcha.Kind.IMAGE:
        alphabet = (string.ascii_uppercase + string.digits).translate(
            {ord(c): None for c in "IO0"}
        )
        length = 5
    elif kind == Captcha.Kind.AUDIO:
        alphabet = string.digits
        length = 8
    else:
        raise ValueError(f"Unknown Captcha kind: {kind}")

    content = "".join([secrets.choice(alphabet) for _ in range(length)])
    metrics.get("desecapi_captcha_content_created").labels(kind).inc()
    return content


class Captcha(ExportModelOperationsMixin("Captcha"), models.Model):
    class Kind(models.TextChoices):
        IMAGE = "image"
        AUDIO = "audio"

    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    created = models.DateTimeField(auto_now_add=True)
    content = models.CharField(max_length=24, default="")
    kind = models.CharField(choices=Kind.choices, default=Kind.IMAGE, max_length=24)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if not self.content:
            self.content = captcha_default_content(self.kind)

    def verify(self, solution: str):
        age = timezone.now() - self.created
        self.delete()
        return (
            str(solution).upper().strip() == self.content  # solution correct
            and age <= settings.CAPTCHA_VALIDITY_PERIOD  # not expired
        )