Feat: API Key management (#1706)

* API Key management backend optimization

* fix lint

* Added code rabbit suggestions

* Initial frontend

* fix check

* Added model

* Added new logic for authentication

* code rabbit suggestion

* fix checks

* backend test fix

* frontend test fix

* added tests fe

* cleanup

* fixes

* Added skeleton for laoding and better gql responses

* fixed cases

* fixed the hook

* fix code

* added testcases

* Update code

* fixed warnings

* Added comments on auth

* Updated code

* updated types and comments on the backend

* Update code (w/o test fixes)

* updated code new suggestions

* fix fe tests

* Added the remaining suggestions

* code rabbit optimizations

* code rabbit optimizations -2

* bug

* Update code

---------

Co-authored-by: Kate Golovanova <[email protected]>
Co-authored-by: Arkadii Yakovets <[email protected]>

Author: 3 people

backend/apps/core/api/ninja.py

@@ -0,0 +1,33 @@
+"""API key authentication class for Django Ninja."""
+
+from http import HTTPStatus
+
+from ninja.errors import HttpError
+from ninja.security import APIKeyHeader
+
+from apps.nest.models.api_key import ApiKey
+
+
+class ApiKeyAuth(APIKeyHeader):
+    """Custom API key authentication class for Ninja."""
+
+    param_name = "X-API-Key"
+
+    def authenticate(self, request, key: str) -> ApiKey:
+        """Authenticate the API key from the request header.
+
+        Args:
+            request: The HTTP request object.
+            key: The API key string from the request header.
+
+        Returns:
+            APIKey: The APIKey object if the key is valid, otherwise None.
+
+        """
+        if not key:
+            raise HttpError(HTTPStatus.UNAUTHORIZED, "Missing API key in 'X-API-Key' header")
+
+        if api_key := ApiKey.authenticate(raw_key=key):
+            return api_key
+
+        raise HttpError(HTTPStatus.UNAUTHORIZED, "Invalid API key")

backend/apps/nest/admin.py

@@ -2,6 +2,7 @@
 
 from django.contrib import admin
 
+from apps.nest.models.api_key import ApiKey
 from apps.nest.models.user import User
 
 
@@ -10,4 +11,25 @@ class UserAdmin(admin.ModelAdmin):
     search_fields = ("email", "username")
 
 
+class ApiKeyAdmin(admin.ModelAdmin):
+    autocomplete_fields = ("user",)
+    list_display = (
+        "name",
+        "user",
+        "uuid",
+        "is_revoked",
+        "expires_at",
+        "created_at",
+        "last_used_at",
+    )
+    list_filter = ("is_revoked",)
+    ordering = ("-created_at",)
+    search_fields = (
+        "name",
+        "uuid",
+        "user__username",
+    )
+
+
+admin.site.register(ApiKey, ApiKeyAdmin)
 admin.site.register(User, UserAdmin)

backend/apps/nest/graphql/__init__.py

@@ -1,3 +1,11 @@
-"""Core User mutations."""
+"""Nest app mutations."""
 
+import strawberry
+
+from .api_key import ApiKeyMutations
 from .user import UserMutations
+
+
+@strawberry.type
+class NestMutations(ApiKeyMutations, UserMutations):
+    """Nest mutations."""

backend/apps/nest/graphql/mutations/__init__.py

@@ -0,0 +1,107 @@
+"""Nest API key GraphQL Mutations."""
+
+import logging
+from datetime import datetime
+from uuid import UUID
+
+import strawberry
+from django.db.utils import IntegrityError
+from django.utils import timezone
+from strawberry.types import Info
+
+from apps.nest.graphql.nodes.api_key import ApiKeyNode
+from apps.nest.graphql.permissions import IsAuthenticated
+from apps.nest.models.api_key import MAX_ACTIVE_KEYS, MAX_WORD_LENGTH, ApiKey
+
+logger = logging.getLogger(__name__)
+
+
+@strawberry.type
+class RevokeApiKeyResult:
+    """Payload for API key revocation result."""
+
+    ok: bool
+    code: str | None = None
+    message: str | None = None
+
+
+@strawberry.type
+class CreateApiKeyResult:
+    """Result of creating an API key."""
+
+    ok: bool
+    api_key: ApiKeyNode | None = None
+    raw_key: str | None = None
+    code: str | None = None
+    message: str | None = None
+
+
+@strawberry.type
+class ApiKeyMutations:
+    """GraphQL mutation class for API keys."""
+
+    @strawberry.mutation(permission_classes=[IsAuthenticated])
+    def create_api_key(self, info: Info, name: str, expires_at: datetime) -> CreateApiKeyResult:
+        """Create a new API key for the authenticated user."""
+        if not name or not name.strip():
+            return CreateApiKeyResult(ok=False, code="INVALID_NAME", message="Name is required")
+
+        if len(name.strip()) > MAX_WORD_LENGTH:
+            return CreateApiKeyResult(ok=False, code="INVALID_NAME", message="Name too long")
+
+        if expires_at <= timezone.now():
+            return CreateApiKeyResult(
+                ok=False, code="INVALID_DATE", message="Expiry date must be in future"
+            )
+
+        try:
+            if not (
+                result := ApiKey.create(
+                    expires_at=expires_at,
+                    name=name,
+                    user=info.context.request.user,
+                )
+            ):
+                return CreateApiKeyResult(
+                    ok=False,
+                    code="LIMIT_REACHED",
+                    message=f"You can have at most {MAX_ACTIVE_KEYS} active API keys.",
+                )
+
+            instance, raw_key = result
+            return CreateApiKeyResult(
+                ok=True,
+                api_key=instance,
+                raw_key=raw_key,
+                code="SUCCESS",
+                message="API key created successfully.",
+            )
+        except IntegrityError as err:
+            logger.warning("Error creating API key: %s", err)
+            return CreateApiKeyResult(
+                ok=False,
+                code="ERROR",
+                message="Something went wrong.",
+            )
+
+    @strawberry.mutation(permission_classes=[IsAuthenticated])
+    def revoke_api_key(self, info: Info, uuid: UUID) -> RevokeApiKeyResult:
+        """Revoke an API key for the authenticated user."""
+        try:
+            api_key = ApiKey.objects.get(
+                uuid=uuid,
+                user=info.context.request.user,
+            )
+            api_key.is_revoked = True
+            api_key.save(update_fields=("is_revoked", "updated_at"))
+        except ApiKey.DoesNotExist:
+            logger.warning("API Key does not exist")
+            return RevokeApiKeyResult(
+                ok=False,
+                code="NOT_FOUND",
+                message="API key not found.",
+            )
+        else:
+            return RevokeApiKeyResult(
+                ok=True, code="SUCCESS", message="API key revoked successfully."
+            )

backend/apps/nest/graphql/mutations/api_key.py

@@ -4,15 +4,27 @@
 
 import requests
 import strawberry
+from django.contrib.auth import login, logout
 from github import Github
+from strawberry.types import Info
 
 from apps.github.models import User as GithubUser
 from apps.nest.graphql.nodes.user import AuthUserNode
+from apps.nest.graphql.permissions import IsAuthenticated
 from apps.nest.models import User
 
 logger = logging.getLogger(__name__)
 
 
+@strawberry.type
+class LogoutResult:
+    """Payload for logout mutation."""
+
+    ok: bool
+    code: str | None = None
+    message: str | None = None
+
+
 @strawberry.type
 class GitHubAuthResult:
     """Payload for GitHubAuth mutation."""
@@ -25,7 +37,7 @@ class UserMutations:
     """GraphQL mutations related to user."""
 
     @strawberry.mutation
-    def github_auth(self, access_token: str) -> GitHubAuthResult:
+    def github_auth(self, info: Info, access_token: str) -> GitHubAuthResult:
         """Authenticate via GitHub OAuth2."""
         try:
             github = Github(access_token)
@@ -49,8 +61,26 @@ def github_auth(self, access_token: str) -> GitHubAuthResult:
                 username=gh_user.login,
             )
 
+            # Log the user in and attach it to a session.
+            # https://docs.djangoproject.com/en/5.2/topics/auth/default/#django.contrib.auth.login
+            # https://docs.djangoproject.com/en/5.2/topics/http/sessions/
+            login(info.context.request, auth_user)
+
             return GitHubAuthResult(auth_user=auth_user)
 
         except requests.exceptions.RequestException as e:
             logger.warning("GitHub authentication failed: %s", e)
             return GitHubAuthResult(auth_user=None)
+
+    @strawberry.mutation(permission_classes=[IsAuthenticated])
+    def logout_user(self, info: Info) -> LogoutResult:
+        """Logout the current user."""
+        # Log the user out and clear the session.
+        # https://docs.djangoproject.com/en/5.2/topics/auth/default/#django.contrib.auth.logout
+        logout(info.context.request)
+
+        return LogoutResult(
+            ok=True,
+            code="SUCCESS",
+            message="User logged out successfully.",
+        )

backend/apps/nest/graphql/mutations/user.py

@@ -0,0 +1,19 @@
+"""GraphQL node for ApiKey model."""
+
+import strawberry_django
+
+from apps.nest.models.api_key import ApiKey
+
+
+@strawberry_django.type(
+    ApiKey,
+    fields=[
+        "created_at",
+        "is_revoked",
+        "expires_at",
+        "name",
+        "uuid",
+    ],
+)
+class ApiKeyNode:
+    """GraphQL node for API keys."""

backend/apps/nest/graphql/nodes/api_key.py

@@ -0,0 +1,21 @@
+"""GraphQL permissions classes for authentication."""
+
+from typing import Any
+
+from graphql import GraphQLError
+from strawberry.permission import BasePermission
+from strawberry.types import Info
+
+
+class IsAuthenticated(BasePermission):
+    """Permission class to check if the user is authenticated."""
+
+    message = "You must be logged in to perform this action."
+
+    def has_permission(self, source, info: Info, **kwargs) -> bool:
+        """Check if the user is authenticated."""
+        return info.context.request.user.is_authenticated
+
+    def on_unauthorized(self) -> Any:
+        """Handle unauthorized access."""
+        return GraphQLError(self.message, extensions={"code": "UNAUTHORIZED"})

backend/apps/nest/graphql/permissions.py

@@ -0,0 +1,8 @@
+import strawberry
+
+from apps.nest.graphql.queries.api_key import ApiKeyQueries
+
+
+@strawberry.type
+class NestQuery(ApiKeyQueries):
+    """Nest query."""

backend/apps/nest/graphql/queries/__init__.py

@@ -0,0 +1,30 @@
+"""GraphQL queries for API keys."""
+
+import strawberry
+from strawberry.types import Info
+
+from apps.nest.graphql.nodes.api_key import ApiKeyNode
+from apps.nest.graphql.permissions import IsAuthenticated
+
+
+@strawberry.type
+class ApiKeyQueries:
+    """GraphQL query class for retrieving API keys."""
+
+    @strawberry.field(permission_classes=[IsAuthenticated])
+    def active_api_key_count(self, info: Info) -> int:
+        """Return count of active API keys for user."""
+        return info.context.request.user.active_api_keys.count()
+
+    @strawberry.field(permission_classes=[IsAuthenticated])
+    def api_keys(self, info: Info) -> list[ApiKeyNode]:
+        """Resolve API keys for the authenticated user.
+
+        Args:
+            info: GraphQL resolver context.
+
+        Returns:
+            list[ApiKeyNode]: List of API keys associated with the authenticated user.
+
+        """
+        return info.context.request.user.active_api_keys.order_by("-created_at")