authentik.blueprints.v1.exporter

  1"""Blueprint exporter"""
  2
  3from collections.abc import Iterable
  4from uuid import UUID
  5
  6from django.apps import apps
  7from django.contrib.auth import get_user_model
  8from django.db.models import Model, Q, QuerySet
  9from django.utils.timezone import now
 10from django.utils.translation import gettext as _
 11from yaml import dump
 12
 13from authentik.blueprints.v1.common import (
 14    Blueprint,
 15    BlueprintDumper,
 16    BlueprintEntry,
 17    BlueprintMetadata,
 18)
 19from authentik.blueprints.v1.importer import is_model_allowed
 20from authentik.blueprints.v1.labels import LABEL_AUTHENTIK_GENERATED
 21from authentik.events.models import Event
 22from authentik.flows.models import Flow, FlowStageBinding, Stage
 23from authentik.policies.models import Policy, PolicyBinding
 24from authentik.stages.prompt.models import PromptStage
 25
 26
 27class Exporter:
 28    """Export flow with attached stages into yaml"""
 29
 30    excluded_models: list[type[Model]] = []
 31
 32    def __init__(self):
 33        self.excluded_models = [
 34            Event,
 35        ]
 36
 37    def get_entries(self) -> Iterable[BlueprintEntry]:
 38        """Get blueprint entries"""
 39        for model in apps.get_models():
 40            if not is_model_allowed(model):
 41                continue
 42            if model in self.excluded_models:
 43                continue
 44            for obj in self.get_model_instances(model):
 45                yield BlueprintEntry.from_model(obj)
 46
 47    def get_model_instances(self, model: type[Model]) -> QuerySet:
 48        """Return a queryset for `model`. Can be used to filter some
 49        objects on some models"""
 50        if model == get_user_model():
 51            return model.objects.exclude_anonymous()
 52        return model.objects.all()
 53
 54    def _pre_export(self, blueprint: Blueprint):
 55        """Hook to run anything pre-export"""
 56
 57    def export(self) -> Blueprint:
 58        """Create a list of all objects and create a blueprint"""
 59        blueprint = Blueprint()
 60        self._pre_export(blueprint)
 61        blueprint.metadata = BlueprintMetadata(
 62            name=_("authentik Export - {date}".format_map({"date": str(now())})),
 63            labels={
 64                LABEL_AUTHENTIK_GENERATED: "true",
 65            },
 66        )
 67        blueprint.entries = list(self.get_entries())
 68        return blueprint
 69
 70    def export_to_string(self) -> str:
 71        """Call export and convert it to yaml"""
 72        blueprint = self.export()
 73        return dump(blueprint, Dumper=BlueprintDumper)
 74
 75
 76class FlowExporter(Exporter):
 77    """Exporter customized to only return objects related to `flow`"""
 78
 79    flow: Flow
 80    with_policies: bool
 81    with_stage_prompts: bool
 82
 83    pbm_uuids: list[UUID]
 84
 85    def __init__(self, flow: Flow):
 86        super().__init__()
 87        self.flow = flow
 88        self.with_policies = True
 89        self.with_stage_prompts = True
 90
 91    def _pre_export(self, blueprint: Blueprint):
 92        if not self.with_policies:
 93            return
 94        self.pbm_uuids = [self.flow.pbm_uuid]
 95        self.pbm_uuids += FlowStageBinding.objects.filter(target=self.flow).values_list(
 96            "pbm_uuid", flat=True
 97        )
 98
 99    def walk_stages(self) -> Iterable[BlueprintEntry]:
100        """Convert all stages attached to self.flow into BlueprintEntry objects"""
101        stages = Stage.objects.filter(flow=self.flow).select_related().select_subclasses()
102        for stage in stages:
103            if isinstance(stage, PromptStage):
104                pass
105            yield BlueprintEntry.from_model(stage, "name")
106
107    def walk_stage_bindings(self) -> Iterable[BlueprintEntry]:
108        """Convert all bindings attached to self.flow into BlueprintEntry objects"""
109        bindings = FlowStageBinding.objects.filter(target=self.flow).select_related()
110        for binding in bindings:
111            yield BlueprintEntry.from_model(binding, "target", "stage", "order")
112
113    def walk_policies(self) -> Iterable[BlueprintEntry]:
114        """Walk over all policies. This is done at the beginning of the export for stages that have
115        a direct foreign key to a policy."""
116        # Special case for PromptStage as that has a direct M2M to policy, we have to ensure
117        # all policies referenced in there we also include here
118        prompt_stages = PromptStage.objects.filter(flow=self.flow).values_list("pk", flat=True)
119        query = Q(bindings__in=self.pbm_uuids) | Q(promptstage__in=prompt_stages)
120        policies = Policy.objects.filter(query).select_related()
121        for policy in policies:
122            yield BlueprintEntry.from_model(policy)
123
124    def walk_policy_bindings(self) -> Iterable[BlueprintEntry]:
125        """Walk over all policybindings relative to us. This is run at the end of the export, as
126        we are sure all objects exist now."""
127        bindings = PolicyBinding.objects.filter(target__in=self.pbm_uuids).select_related()
128        for binding in bindings:
129            yield BlueprintEntry.from_model(binding, "policy", "target", "order")
130
131    def walk_stage_prompts(self) -> Iterable[BlueprintEntry]:
132        """Walk over all prompts associated with any PromptStages"""
133        prompt_stages = PromptStage.objects.filter(flow=self.flow)
134        for stage in prompt_stages:
135            for prompt in stage.fields.all():
136                yield BlueprintEntry.from_model(prompt)
137
138    def get_entries(self) -> Iterable[BlueprintEntry]:
139        entries = []
140        entries.append(BlueprintEntry.from_model(self.flow, "slug"))
141        if self.with_stage_prompts:
142            entries.extend(self.walk_stage_prompts())
143        if self.with_policies:
144            entries.extend(self.walk_policies())
145        entries.extend(self.walk_stages())
146        entries.extend(self.walk_stage_bindings())
147        if self.with_policies:
148            entries.extend(self.walk_policy_bindings())
149        return entries