authentik.stages.prompt.api

  1"""Prompt Stage API Views"""
  2
  3from drf_spectacular.utils import extend_schema
  4from rest_framework.decorators import action
  5from rest_framework.request import Request
  6from rest_framework.response import Response
  7from rest_framework.serializers import CharField
  8from rest_framework.validators import UniqueValidator
  9from rest_framework.viewsets import ModelViewSet
 10
 11from authentik.core.api.used_by import UsedByMixin
 12from authentik.core.api.utils import ModelSerializer
 13from authentik.core.expression.exceptions import PropertyMappingExpressionException
 14from authentik.flows.api.stages import StageSerializer
 15from authentik.flows.challenge import HttpChallengeResponse
 16from authentik.flows.planner import FlowPlan
 17from authentik.flows.views.executor import FlowExecutorView
 18from authentik.lib.generators import generate_id
 19from authentik.lib.utils.errors import exception_to_string
 20from authentik.stages.prompt.models import Prompt, PromptStage
 21from authentik.stages.prompt.stage import PromptChallenge, PromptStageView
 22
 23
 24class PromptStageSerializer(StageSerializer):
 25    """PromptStage Serializer"""
 26
 27    name = CharField(validators=[UniqueValidator(queryset=PromptStage.objects.all())])
 28
 29    class Meta:
 30        model = PromptStage
 31        fields = StageSerializer.Meta.fields + [
 32            "fields",
 33            "validation_policies",
 34        ]
 35
 36
 37class PromptStageViewSet(UsedByMixin, ModelViewSet):
 38    """PromptStage Viewset"""
 39
 40    queryset = PromptStage.objects.prefetch_related(
 41        "flow_set",
 42        "fields",
 43        "validation_policies",
 44    ).all()
 45    serializer_class = PromptStageSerializer
 46    filterset_fields = "__all__"
 47    ordering = ["name"]
 48    search_fields = ["name"]
 49
 50
 51class PromptSerializer(ModelSerializer):
 52    """Prompt Serializer"""
 53
 54    prompt_stages_obj = PromptStageSerializer(
 55        source="promptstage_set", many=True, required=False, read_only=True
 56    )
 57
 58    class Meta:
 59        model = Prompt
 60        fields = [
 61            "pk",
 62            "name",
 63            "field_key",
 64            "label",
 65            "type",
 66            "required",
 67            "placeholder",
 68            "initial_value",
 69            "order",
 70            "prompt_stages_obj",
 71            "sub_text",
 72            "placeholder_expression",
 73            "initial_value_expression",
 74        ]
 75
 76
 77class PromptViewSet(UsedByMixin, ModelViewSet):
 78    """Prompt Viewset"""
 79
 80    queryset = Prompt.objects.all().prefetch_related(
 81        "promptstage_set",
 82        "promptstage_set__flow_set",
 83        "promptstage_set__fields",
 84        "promptstage_set__validation_policies",
 85    )
 86    serializer_class = PromptSerializer
 87    ordering = ["field_key"]
 88    filterset_fields = ["field_key", "name", "label", "type", "placeholder"]
 89    search_fields = ["field_key", "name", "label", "type", "placeholder"]
 90
 91    @extend_schema(
 92        request=PromptSerializer,
 93        responses={
 94            200: PromptChallenge,
 95        },
 96    )
 97    @action(detail=False, methods=["POST"])
 98    def preview(self, request: Request) -> Response:
 99        """Preview a prompt as a challenge, just like a flow would receive"""
100        # Remove a couple things from the request, the serializer will fail on these
101        # when previewing an existing prompt
102        # and since we don't plan to save from this, set a random name and remove the stage
103        request.data["name"] = generate_id()
104        request.data.pop("promptstage_set", None)
105        # Validate data, same as a normal edit/create request
106        prompt = PromptSerializer(data=request.data)
107        prompt.is_valid(raise_exception=True)
108        # Convert serializer to prompt instance
109        prompt_model = Prompt(**prompt.validated_data)
110        # Convert to field challenge
111        try:
112            fields = PromptStageView(
113                FlowExecutorView(
114                    plan=FlowPlan(""),
115                    request=request._request,
116                ),
117                request=request._request,
118            ).get_prompt_challenge_fields([prompt_model], {}, dry_run=True)
119        except PropertyMappingExpressionException as exc:
120            return Response(
121                {
122                    "non_field_errors": [
123                        exception_to_string(exc.exc),
124                    ]
125                },
126                status=400,
127            )
128        challenge = PromptChallenge(
129            data={
130                "fields": fields,
131            },
132        )
133        challenge.is_valid()
134        return HttpChallengeResponse(challenge)