authentik.endpoints.connectors.agent.tests.test_stage

  1from hashlib import sha256
  2from json import loads
  3from unittest.mock import PropertyMock, patch
  4
  5from django.urls import reverse
  6from jwt import encode
  7
  8from authentik.core.tests.utils import create_test_cert, create_test_flow
  9from authentik.endpoints.connectors.agent.models import (
 10    AgentConnector,
 11    AgentDeviceConnection,
 12    DeviceAuthenticationToken,
 13    DeviceToken,
 14    EnrollmentToken,
 15)
 16from authentik.endpoints.connectors.agent.stage import (
 17    PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE,
 18    PLAN_CONTEXT_DEVICE_AUTH_TOKEN,
 19)
 20from authentik.endpoints.models import Device, EndpointStage, StageMode
 21from authentik.flows.models import FlowStageBinding
 22from authentik.flows.planner import PLAN_CONTEXT_DEVICE
 23from authentik.flows.tests import FlowTestCase
 24from authentik.lib.generators import generate_id
 25
 26
 27class TestEndpointStage(FlowTestCase):
 28
 29    def setUp(self):
 30        cert = create_test_cert()
 31        self.connector = AgentConnector.objects.create(name=generate_id(), challenge_key=cert)
 32        self.token = EnrollmentToken.objects.create(name=generate_id(), connector=self.connector)
 33        self.device = Device.objects.create(
 34            identifier=generate_id(),
 35        )
 36        self.connection = AgentDeviceConnection.objects.create(
 37            device=self.device,
 38            connector=self.connector,
 39        )
 40        self.device_token = DeviceToken.objects.create(
 41            device=self.connection,
 42            key=generate_id(),
 43        )
 44        self.device_auth_token = DeviceAuthenticationToken.objects.create(
 45            device=self.device,
 46            device_token=self.device_token,
 47            connector=self.connector,
 48        )
 49
 50    def test_endpoint_stage(self):
 51        flow = create_test_flow()
 52        stage = EndpointStage.objects.create(connector=self.connector)
 53        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
 54
 55        res = self.client.get(
 56            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
 57        )
 58        self.assertEqual(res.status_code, 200)
 59        self.assertStageResponse(res, flow=flow, component="ak-stage-endpoint-agent")
 60
 61        challenge = loads(res.content.decode())["challenge"]
 62
 63        DeviceToken.objects.create(
 64            device=self.connection,
 65            key=generate_id(),
 66        )
 67
 68        response = encode(
 69            {
 70                "iss": self.device.identifier,
 71                "atc": challenge,
 72                "aud": "goauthentik.io/platform/endpoint",
 73            },
 74            key=self.device_token.key,
 75            algorithm="HS512",
 76        )
 77
 78        with self.assertFlowFinishes() as plan:
 79            res = self.client.post(
 80                reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
 81                data={"component": "ak-stage-endpoint-agent", "response": response},
 82            )
 83            self.assertStageRedirects(res, reverse("authentik_core:root-redirect"))
 84        plan = plan()
 85        self.assertNotIn(PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE, plan.context)
 86        self.assertEqual(plan.context[PLAN_CONTEXT_DEVICE], self.device)
 87
 88    def test_endpoint_stage_optional(self):
 89        flow = create_test_flow()
 90        stage = EndpointStage.objects.create(connector=self.connector, mode=StageMode.OPTIONAL)
 91        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
 92
 93        res = self.client.get(
 94            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
 95        )
 96        self.assertEqual(res.status_code, 200)
 97        self.assertStageResponse(res, flow=flow, component="ak-stage-endpoint-agent")
 98
 99        challenge = loads(res.content.decode())["challenge"]
100
101        response = encode(
102            {
103                "iss": self.device.identifier,
104                "atc": challenge,
105                "aud": "goauthentik.io/platform/endpoint",
106            },
107            key=self.device_token.key,
108            algorithm="HS512",
109        )
110
111        with self.assertFlowFinishes() as plan:
112            res = self.client.post(
113                reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
114                data={"component": "ak-stage-endpoint-agent", "response": response},
115            )
116            self.assertStageRedirects(res, reverse("authentik_core:root-redirect"))
117        plan = plan()
118        self.assertNotIn(PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE, plan.context)
119        self.assertEqual(plan.context[PLAN_CONTEXT_DEVICE], self.device)
120
121    def test_endpoint_stage_optional_none(self):
122        flow = create_test_flow()
123        stage = EndpointStage.objects.create(connector=self.connector, mode=StageMode.OPTIONAL)
124        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
125
126        res = self.client.get(
127            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
128        )
129        self.assertEqual(res.status_code, 200)
130        self.assertStageResponse(res, flow=flow, component="ak-stage-endpoint-agent")
131
132        with self.assertFlowFinishes() as plan:
133            res = self.client.post(
134                reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
135                data={"component": "ak-stage-endpoint-agent", "response": None},
136            )
137            self.assertStageRedirects(res, reverse("authentik_core:root-redirect"))
138        plan = plan()
139        self.assertNotIn(PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE, plan.context)
140        self.assertNotIn(PLAN_CONTEXT_DEVICE, plan.context)
141
142    def test_endpoint_stage_optional_invalid(self):
143        flow = create_test_flow()
144        stage = EndpointStage.objects.create(connector=self.connector, mode=StageMode.OPTIONAL)
145        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
146
147        res = self.client.get(
148            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
149        )
150        self.assertEqual(res.status_code, 200)
151        self.assertStageResponse(res, flow=flow, component="ak-stage-endpoint-agent")
152
153        with self.assertFlowFinishes() as plan:
154            res = self.client.post(
155                reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
156                data={"component": "ak-stage-endpoint-agent", "response": generate_id()},
157            )
158            self.assertStageRedirects(res, reverse("authentik_core:root-redirect"))
159        plan = plan()
160        self.assertNotIn(PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE, plan.context)
161        self.assertNotIn(PLAN_CONTEXT_DEVICE, plan.context)
162
163    def test_endpoint_stage_required_none(self):
164        flow = create_test_flow()
165        stage = EndpointStage.objects.create(connector=self.connector, mode=StageMode.REQUIRED)
166        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
167
168        res = self.client.get(
169            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
170        )
171        self.assertEqual(res.status_code, 200)
172        self.assertStageResponse(res, flow=flow, component="ak-stage-endpoint-agent")
173
174        res = self.client.post(
175            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
176            data={"component": "ak-stage-endpoint-agent", "response": None},
177        )
178        self.assertStageResponse(
179            res,
180            flow=flow,
181            component="ak-stage-access-denied",
182            error_message="Invalid challenge response",
183        )
184
185    def test_endpoint_stage_required_invalid(self):
186        flow = create_test_flow()
187        stage = EndpointStage.objects.create(connector=self.connector, mode=StageMode.REQUIRED)
188        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
189
190        res = self.client.get(
191            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
192        )
193        self.assertEqual(res.status_code, 200)
194        self.assertStageResponse(res, flow=flow, component="ak-stage-endpoint-agent")
195
196        res = self.client.post(
197            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
198            data={"component": "ak-stage-endpoint-agent", "response": generate_id()},
199        )
200        self.assertStageResponse(
201            res,
202            flow=flow,
203            component="ak-stage-endpoint-agent",
204            response_errors={
205                "response": [{"string": "Invalid challenge response", "code": "invalid"}]
206            },
207        )
208
209    def test_endpoint_stage_ia_dth(self):
210        """Test with DTH"""
211        flow = create_test_flow()
212        stage = EndpointStage.objects.create(connector=self.connector)
213        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
214
215        # Send an "invalid" request first, to populate the flow plan
216        res = self.client.get(
217            reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
218        )
219        plan = self.get_flow_plan()
220        plan.context[PLAN_CONTEXT_DEVICE_AUTH_TOKEN] = DeviceAuthenticationToken.objects.get(
221            pk=self.device_auth_token.pk
222        )
223        self.set_flow_plan(plan)
224
225        with self.assertFlowFinishes() as plan:
226            res = self.client.get(
227                reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
228                HTTP_X_AUTHENTIK_PLATFORM_AUTH_DTH=sha256(
229                    self.device_token.key.encode()
230                ).hexdigest(),
231            )
232            self.assertStageRedirects(res, reverse("authentik_core:root-redirect"))
233        plan = plan()
234        self.assertNotIn(PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE, plan.context)
235        self.assertEqual(plan.context[PLAN_CONTEXT_DEVICE], self.device)
236
237    def test_endpoint_stage_connector_no_stage_optional(self):
238        flow = create_test_flow()
239        stage = EndpointStage.objects.create(connector=self.connector, mode=StageMode.OPTIONAL)
240        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
241
242        with patch(
243            "authentik.endpoints.connectors.agent.models.AgentConnector.stage",
244            PropertyMock(return_value=None),
245        ):
246            with self.assertFlowFinishes() as plan:
247                res = self.client.get(
248                    reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
249                )
250                self.assertStageRedirects(res, reverse("authentik_core:root-redirect"))
251            plan = plan()
252            self.assertNotIn(PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE, plan.context)
253            self.assertNotIn(PLAN_CONTEXT_DEVICE, plan.context)
254
255    def test_endpoint_stage_connector_no_stage_required(self):
256        flow = create_test_flow()
257        stage = EndpointStage.objects.create(connector=self.connector, mode=StageMode.REQUIRED)
258        FlowStageBinding.objects.create(stage=stage, target=flow, order=0)
259
260        with patch(
261            "authentik.endpoints.connectors.agent.models.AgentConnector.stage",
262            PropertyMock(return_value=None),
263        ):
264            with self.assertFlowFinishes() as plan:
265                res = self.client.get(
266                    reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}),
267                )
268                self.assertStageResponse(
269                    res,
270                    component="ak-stage-access-denied",
271                    error_message="Invalid stage configuration",
272                )
273            plan = plan()
274            self.assertNotIn(PLAN_CONTEXT_AGENT_ENDPOINT_CHALLENGE, plan.context)
275            self.assertNotIn(PLAN_CONTEXT_DEVICE, plan.context)