diff --git a/arkindex/ponos/authentication.py b/arkindex/ponos/authentication.py index 8eaa6bf008d9abd3500dbf5bdf3b538fb62e9bfc..152619cb67810be0cd1fa5296134d4059ae4044d 100644 --- a/arkindex/ponos/authentication.py +++ b/arkindex/ponos/authentication.py @@ -1,7 +1,9 @@ +from drf_spectacular.authentication import TokenScheme from drf_spectacular.contrib.rest_framework_simplejwt import SimpleJWTScheme +from rest_framework.authentication import TokenAuthentication from rest_framework.exceptions import AuthenticationFailed -from arkindex.ponos.models import Agent +from arkindex.ponos.models import Agent, Task from rest_framework_simplejwt.authentication import JWTAuthentication from rest_framework_simplejwt.exceptions import InvalidToken from rest_framework_simplejwt.settings import api_settings @@ -68,3 +70,35 @@ class AgentAuthenticationExtension(SimpleJWTScheme): target_class = "arkindex.ponos.authentication.AgentAuthentication" name = "agentAuth" + + +class TaskAuthentication(TokenAuthentication): + keyword = 'Ponos' + model = Task + + def authenticate_credentials(self, key): + try: + task = Task.objects.select_related('workflow__process__creator').get(token=key) + except Task.DoesNotExist: + # Same error message as the standard TokenAuthentication + raise AuthenticationFailed('Invalid token.') + + if not task.workflow.process: + raise AuthenticationFailed('Task has no process.') + + user = task.workflow.process.creator + if not user or not user.is_active: + # Same error message as the standard TokenAuthentication + raise AuthenticationFailed('User inactive or deleted.') + + return (user, task) + + +class TaskAuthenticationExtension(TokenScheme): + target_class = "arkindex.ponos.authentication.TaskAuthentication" + name = "taskAuth" + # The TokenScheme has a priority of -1 and matches both TokenAuthentication and its subclasses; + # we set the priority to a higher number to make this extension match first, and disable + # subclass matching so that this only applies to the TaskAuthentication. + priority = 0 + match_subclasses = False diff --git a/arkindex/project/settings.py b/arkindex/project/settings.py index b5e354e42b09227551003cbe2c025cc5d00a1013..820a6c4de196c0b75dbfc0f1ba3b43e47e9e5113 100644 --- a/arkindex/project/settings.py +++ b/arkindex/project/settings.py @@ -202,6 +202,7 @@ REST_FRAMEWORK = { 'rest_framework.authentication.SessionAuthentication', 'rest_framework.authentication.TokenAuthentication', 'arkindex.ponos.authentication.AgentAuthentication', + 'arkindex.ponos.authentication.TaskAuthentication', ), 'DEFAULT_PAGINATION_CLASS': 'arkindex.project.pagination.PageNumberPagination', 'DEFAULT_SCHEMA_CLASS': 'arkindex.project.openapi.AutoSchema', diff --git a/arkindex/project/tests/openapi/test_schema.py b/arkindex/project/tests/openapi/test_schema.py index 1dd19e06d262faafdbf18951997c2f46e22a2b93..ead901bae65971f8a0bcaa38072f54f8d7149c90 100644 --- a/arkindex/project/tests/openapi/test_schema.py +++ b/arkindex/project/tests/openapi/test_schema.py @@ -56,6 +56,7 @@ class TestAutoSchema(TestCase): {'cookieAuth': []}, {'tokenAuth': []}, {'agentAuth': []}, + {'taskAuth': []}, # Allows no authentication too {}, ],