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

from base64 import b64encode

from captcha.audio import AudioCaptcha
from captcha.image import ImageCaptcha
from rest_framework import serializers

from api import settings
from desecapi.models import Captcha


class CaptchaSerializer(serializers.ModelSerializer):
    challenge = serializers.SerializerMethodField()

    class Meta:
        model = Captcha
        fields = (
            ("id", "challenge", "kind")
            if not settings.DEBUG
            else ("id", "challenge", "kind", "content")
        )

    def get_challenge(self, obj: Captcha):
        # TODO Does this need to be stored in the object instance, in case this method gets called twice?
        if obj.kind == Captcha.Kind.IMAGE:
            challenge = ImageCaptcha().generate(obj.content).getvalue()
        elif obj.kind == Captcha.Kind.AUDIO:
            challenge = AudioCaptcha().generate(obj.content)
        else:
            raise ValueError(f"Unknown captcha type {obj.kind}")
        return b64encode(challenge)


class CaptchaSolutionSerializer(serializers.Serializer):
    id = serializers.PrimaryKeyRelatedField(
        queryset=Captcha.objects.all(),
        error_messages={"does_not_exist": "CAPTCHA does not exist."},
    )
    solution = serializers.CharField(write_only=True, required=True)

    def validate(self, attrs):
        captcha = attrs["id"]  # Note that this already is the Captcha object
        if not captcha.verify(attrs["solution"]):
            raise serializers.ValidationError(
                "CAPTCHA could not be validated. Please obtain a new one and try again."
            )

        return attrs