🔑 feat: revocable API keys with scopes

estela-api/api/authentication.py

@@ -0,0 +1,81 @@
+import hashlib
+import secrets
+
+from django.utils import timezone
+from rest_framework import authentication, exceptions
+
+from core.models import ApiKey
+
+KEYWORD = "Token"
+LAST_USED_RESOLUTION = 300
+
+
+def generate_key():
+    """Returns (plaintext, prefix, hash). The plaintext is shown to the user once."""
+    plaintext = ApiKey.KEY_PREFIX + secrets.token_urlsafe(32)
+    return plaintext, plaintext[: ApiKey.PREFIX_LENGTH], hash_key(plaintext)
+
+
+def hash_key(plaintext):
+    return hashlib.sha256(plaintext.encode()).hexdigest()
+
+
+class ApiKeyAuthentication(authentication.BaseAuthentication):
+    """Authenticates `Authorization: Token estela_...`.
+
+    Anything that is not an estela key is passed on untouched, so the DRF tokens
+    that the web and the CLI use today keep working.
+    """
+
+    def authenticate(self, request):
+        header = authentication.get_authorization_header(request).split()
+        if len(header) != 2 or header[0].decode().lower() != KEYWORD.lower():
+            return None
+
+        try:
+            plaintext = header[1].decode()
+        except UnicodeError:
+            return None
+
+        if not plaintext.startswith(ApiKey.KEY_PREFIX):
+            return None
+
+        try:
+            api_key = ApiKey.objects.select_related("user").get(
+                key_hash=hash_key(plaintext)
+            )
+        except ApiKey.DoesNotExist:
+            raise exceptions.AuthenticationFailed("Invalid API key.")
+
+        if api_key.revoked:
+            raise exceptions.AuthenticationFailed("This API key was revoked.")
+
+        if api_key.expired:
+            raise exceptions.AuthenticationFailed(
+                "This API key expired on {}.".format(
+                    api_key.expires_at.date().isoformat()
+                )
+            )
+
+        if not api_key.user.is_active:
+            raise exceptions.AuthenticationFailed("User inactive or deleted.")
+
+        self.touch(api_key)
+        return (api_key.user, api_key)
+
+    def touch(self, api_key):
+        """Records usage, at most once every LAST_USED_RESOLUTION seconds.
+
+        Airflow polls job status in a loop for hours, so writing on every request
+        would be a write per call.
+        """
+        now = timezone.now()
+        if (
+            api_key.last_used_at
+            and (now - api_key.last_used_at).total_seconds() < LAST_USED_RESOLUTION
+        ):
+            return
+        ApiKey.objects.filter(pk=api_key.pk).update(last_used_at=now)
+
+    def authenticate_header(self, request):
+        return KEYWORD

estela-api/api/mixins.py

@@ -4,7 +4,8 @@
 from rest_framework.pagination import PageNumberPagination
 from rest_framework.permissions import IsAuthenticated
 
-from api.permissions import IsAdminOrReadOnly, IsProjectUser
+from api.authentication import ApiKeyAuthentication
+from api.permissions import HasApiKeyScope, IsAdminOrReadOnly, IsProjectUser
 from core.models import Notification, Activity
 
 
@@ -17,8 +18,13 @@ class APIPageNumberPagination(PageNumberPagination):
 class BaseViewSet(viewsets.GenericViewSet):
     """A custom viewset that contains reusable customized settings."""
 
-    authentication_classes = [TokenAuthentication]
-    permission_classes = [IsAuthenticated, IsProjectUser, IsAdminOrReadOnly]
+    authentication_classes = [ApiKeyAuthentication, TokenAuthentication]
+    permission_classes = [
+        IsAuthenticated,
+        HasApiKeyScope,
+        IsProjectUser,
+        IsAdminOrReadOnly,
+    ]
     pagination_class = APIPageNumberPagination
 
 

estela-api/api/permissions.py

@@ -3,7 +3,38 @@
 from rest_framework.permissions import BasePermission, SAFE_METHODS
 from django.contrib.auth.models import User
 
-from core.models import Project, Permission
+from core.models import ApiKey, Project, Permission
+
+
+class HasApiKeyScope(BasePermission):
+    """Narrows what an API key may do. Sessions are unaffected.
+
+    A key reads whatever its owner can read, unless the view asks for a scope.
+    Writing always needs a scope, so a view that declares none is session-only.
+    """
+
+    message = "This API key does not have the required scope."
+
+    def has_permission(self, request, view):
+        api_key = request.auth
+        if not isinstance(api_key, ApiKey):
+            return True
+
+        if request.method in SAFE_METHODS:
+            required = getattr(view, "api_key_read_scope", None)
+            return required is None or api_key.has_scope(required)
+
+        required = getattr(view, "api_key_write_scope", None)
+        return required is not None and api_key.has_scope(required)
+
+
+class IsSessionAuthenticated(BasePermission):
+    """Blocks API keys outright, for anything only a person should do."""
+
+    message = "This action requires a login session, not an API key."
+
+    def has_permission(self, request, view):
+        return not isinstance(request.auth, ApiKey)
 
 
 class IsProjectUser(BasePermission):

estela-api/api/serializers/apikey.py

@@ -0,0 +1,57 @@
+from django.conf import settings
+from rest_framework import serializers
+
+from core.models import ApiKey
+
+
+class ApiKeySerializer(serializers.ModelSerializer):
+    """What a key looks like once created. Never carries the secret."""
+
+    # Declared so the schema says "array of string"; a bare JSONField documents as a string.
+    scopes = serializers.ListField(child=serializers.CharField(), read_only=True)
+
+    class Meta:
+        model = ApiKey
+        fields = [
+            "id",
+            "name",
+            "prefix",
+            "scopes",
+            "created",
+            "last_used_at",
+            "expires_at",
+        ]
+
+
+class ApiKeyCreateSerializer(serializers.ModelSerializer):
+    scopes = serializers.ListField(
+        child=serializers.ChoiceField(choices=ApiKey.SCOPES),
+        required=False,
+        default=list,
+        help_text="Extra permissions. Empty means read-only.",
+    )
+
+    expires_in_days = serializers.ChoiceField(
+        choices=settings.API_KEY_EXPIRY_CHOICES,
+        required=False,
+        write_only=True,
+        help_text="How long the key lasts. Omit it to get the deployment's default.",
+    )
+
+    class Meta:
+        model = ApiKey
+        fields = ["name", "scopes", "expires_in_days"]
+
+    def validate_scopes(self, value):
+        return sorted(set(value))
+
+
+class ApiKeyCreateResponseSerializer(ApiKeySerializer):
+    """Schema only: tells Swagger that creation also returns the key itself."""
+
+    key = serializers.CharField(
+        read_only=True, help_text="The key itself. It is not stored and never shown again."
+    )
+
+    class Meta(ApiKeySerializer.Meta):
+        fields = ApiKeySerializer.Meta.fields + ["key"]

estela-api/api/urls.py

@@ -12,6 +12,7 @@
     notification as notification_views,
     proxyprovider as proxyprovider_views,
     metering as metering_views,
+    apikey as apikey_views,
 )
 
 router = routers.DefaultRouter(trailing_slash=False)
@@ -64,6 +65,11 @@
     prefix=r"proxy_provider",
     viewset=proxyprovider_views.ProxyProviderViewSet,
 )
+router.register(
+    prefix=r"account/api-keys",
+    viewset=apikey_views.ApiKeyViewSet,
+    basename="api-key",
+)
 router.register(prefix=r"auth", viewset=auth_views.AuthAPIViewSet, basename="auth")
 router.register(
     prefix=r"auth/profile", viewset=auth_views.UserProfileViewSet, basename="profile"

estela-api/api/utils.py

@@ -4,11 +4,12 @@
 
 from django.conf import settings
 import redis
+from rest_framework.authtoken.models import Token
 
 from api import errors
 from api.exceptions import DataBaseError
 from config.job_manager import spiderdata_db_client
-from core.models import SpiderJobEnvVar
+from core.models import Permission, SpiderJobEnvVar
 
 
 def update_env_vars(instance, env_vars, level="project", delete=True):
@@ -309,3 +310,25 @@ def get_job_spiderdata_corpus_sizes_bytes(job) -> Optional[Tuple[int, int, int]]
         spiderdata_db_client.get_dataset_size(pid, get_collection_name(job, "logs"))
     )
     return (items_size, requests_size, logs_size)
+
+
+def get_project_owner_token(job):
+    """The project owner's DRF token, used by job containers to call back."""
+    permission = job.spider.project.permission_set.filter(
+        permission=Permission.OWNER_PERMISSION
+    ).first()
+    if not permission:
+        return None
+    token, _ = Token.objects.get_or_create(user=permission.user)
+    return token.key
+
+
+def get_job_callback_token(auth, job):
+    """An API key cannot be forwarded to the container: only its hash is stored.
+
+    So a request authenticated with one falls back to the owner's token, which is
+    what queued jobs already use.
+    """
+    if isinstance(auth, Token):
+        return auth.key
+    return get_project_owner_token(job)

estela-api/api/views/apikey.py

@@ -0,0 +1,67 @@
+from datetime import timedelta
+
+from django.conf import settings
+from django.utils import timezone
+from drf_yasg.utils import swagger_auto_schema
+from rest_framework import mixins, permissions, status, viewsets
+from rest_framework.authentication import TokenAuthentication
+from rest_framework.response import Response
+
+from api.authentication import ApiKeyAuthentication, generate_key
+from api.permissions import IsSessionAuthenticated
+from api.serializers.apikey import (
+    ApiKeyCreateResponseSerializer,
+    ApiKeyCreateSerializer,
+    ApiKeySerializer,
+)
+from core.models import ApiKey
+
+
+class ApiKeyViewSet(viewsets.GenericViewSet, mixins.ListModelMixin):
+    """Manage the caller's own API keys.
+
+    ApiKeyAuthentication is listed so a request made with a key is recognised and
+    then refused by IsSessionAuthenticated: a leaked key must not be able to mint
+    more keys.
+    """
+
+    authentication_classes = [ApiKeyAuthentication, TokenAuthentication]
+    permission_classes = [permissions.IsAuthenticated, IsSessionAuthenticated]
+    serializer_class = ApiKeySerializer
+
+    def get_queryset(self):
+        if self.request is None:
+            return ApiKey.objects.none()
+        return ApiKey.objects.filter(user=self.request.user, revoked_at__isnull=True)
+
+    @swagger_auto_schema(
+        request_body=ApiKeyCreateSerializer,
+        responses={status.HTTP_201_CREATED: ApiKeyCreateResponseSerializer()},
+    )
+    def create(self, request, *args, **kwargs):
+        serializer = ApiKeyCreateSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+
+        days = serializer.validated_data.get(
+            "expires_in_days", settings.API_KEY_DEFAULT_DAYS
+        )
+        plaintext, prefix, key_hash = generate_key()
+        api_key = ApiKey.objects.create(
+            user=request.user,
+            name=serializer.validated_data["name"],
+            scopes=serializer.validated_data["scopes"],
+            prefix=prefix,
+            key_hash=key_hash,
+            expires_at=timezone.now() + timedelta(days=int(days)),
+        )
+
+        data = ApiKeySerializer(api_key).data
+        data["key"] = plaintext
+        return Response(data, status=status.HTTP_201_CREATED)
+
+    @swagger_auto_schema(responses={status.HTTP_204_NO_CONTENT: "Key revoked"})
+    def destroy(self, request, *args, **kwargs):
+        api_key = self.get_object()
+        api_key.revoked_at = timezone.now()
+        api_key.save(update_fields=["revoked_at"])
+        return Response(status=status.HTTP_204_NO_CONTENT)

estela-api/api/views/cronjob.py

@@ -15,7 +15,7 @@
     SpiderCronJobUpdateSerializer,
 )
 from core.cronjob import create_cronjob, disable_cronjob, run_cronjob_once
-from core.models import DataStatus, Spider, SpiderCronJob, Project
+from core.models import ApiKey, DataStatus, Spider, SpiderCronJob, Project
 from core.tiers import DEFAULT_TIER
 
 
@@ -27,6 +27,8 @@ class SpiderCronJobViewSet(
     mixins.UpdateModelMixin,
     mixins.ListModelMixin,
 ):
+    api_key_write_scope = ApiKey.RUN_SCOPE
+
     model_class = SpiderCronJob
     serializer_class = SpiderCronJobSerializer
     lookup_field = "cjid"

estela-api/api/views/deploy.py

@@ -13,7 +13,7 @@
     DeployUpdateSerializer,
 )
 from config.job_manager import credentials, spiderdata_db_client
-from core.models import Deploy, Project
+from core.models import ApiKey, Deploy, Project
 from core.views import launch_deploy_job
 
 
@@ -22,6 +22,8 @@ class DeployViewSet(
     viewsets.ModelViewSet,
     ActionHandlerMixin,
 ):
+    api_key_write_scope = ApiKey.MANAGE_SCOPE
+
     model_class = Deploy
     serializer_class = DeploySerializer
     lookup_field = "did"