From 13ab6a0ddef89b5cc2fe9c94b1853a98cc339e9f Mon Sep 17 00:00:00 2001
From: lihuacai <lihuacai168@gmail.com>
Date: Fri, 17 Nov 2023 02:35:29 +0800
Subject: [PATCH] feat(core): Update settings, routes, and auth with JWT
 authentication

In settings.py, a new token refresh schema was added to improve our authentication methods. In urls.py, imports were rearranged and new routers were added to handle JWT token input and output; the handling of JWT-related exceptions was also included. In auth.py, extended functionalities were added to handle token generation, token refreshing and retrieval of response schema. This is essential in establishing more secure way to transmit information between parties and prevent unauthorized access.
---
 apidemo/settings.py |  1 +
 apidemo/urls.py     | 31 +++++++++++++++++++++++++------
 core/auth.py        | 39 +++++++++++++++++++++++++++++++++------
 3 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/apidemo/settings.py b/apidemo/settings.py
index 796e357..162c9b8 100644
--- a/apidemo/settings.py
+++ b/apidemo/settings.py
@@ -244,6 +244,7 @@
     # For Controller Schemas
     # FOR OBTAIN PAIR
     'TOKEN_OBTAIN_PAIR_INPUT_SCHEMA': "core.auth.MyTokenObtainPairInputSchema",
+    'TOKEN_OBTAIN_PAIR_REFRESH_INPUT_SCHEMA': "core.auth.MyTokenRefreshInputSchema",
     # 'TOKEN_OBTAIN_PAIR_REFRESH_INPUT_SCHEMA': "ninja_jwt.schema.TokenRefreshInputSchema",
     # # FOR SLIDING TOKEN
     # 'TOKEN_OBTAIN_SLIDING_INPUT_SCHEMA': "ninja_jwt.schema.TokenObtainSlidingInputSchema",
diff --git a/apidemo/urls.py b/apidemo/urls.py
index be08b72..5fc5a26 100644
--- a/apidemo/urls.py
+++ b/apidemo/urls.py
@@ -13,21 +13,40 @@
     1. Import the include() function: from django.urls import include, path
     2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
 """
+from django.contrib import admin
+from django.urls import path
+from ninja import File, NinjaAPI
+from ninja_extra import exceptions as extra_exceptions
 from ninja_jwt.routers.obtain import obtain_pair_router
 
 from employee.views import router as employee_router
 
+api_v1 = NinjaAPI(version="1.0.0")
 
-from django.contrib import admin
-from django.urls import path
-from ninja import NinjaAPI, File
+api_v1.add_router("/employee/", employee_router)
+api_v1.add_router("/token", tags=["Auth"], router=obtain_pair_router)
+
+
+def obtain_token_exception_handler(request, exc):
+    headers = {}
+
+    if isinstance(exc, extra_exceptions.APIException):
+        data = {
+            "message": exc.detail.get("detail", str(exc)),
+            "success": False,
+            "data": None,
+        }
+    else:
+        data = {"message": exc.detail, "success": False, "data": None}
 
+    response = api_v1.create_response(request, data, status=exc.status_code)
+    for k, v in headers.items():
+        response.setdefault(k, v)
 
-api_v1 = NinjaAPI(version='1.0.0')
+    return response
 
-api_v1.add_router('/employee/', employee_router)
-api_v1.add_router('/token', tags=['Auth'], router=obtain_pair_router)
 
+api_v1.exception_handler(extra_exceptions.APIException)(obtain_token_exception_handler)
 
 urlpatterns = [
     path("admin/", admin.site.urls),
diff --git a/core/auth.py b/core/auth.py
index 23681b7..48a18ef 100644
--- a/core/auth.py
+++ b/core/auth.py
@@ -1,11 +1,15 @@
 # !/usr/bin/python3
 # -*- coding: utf-8 -*-
-from typing import Dict, Type
+from typing import Dict, Optional, Type, Union
 
 from django.contrib.auth import get_user_model
 from django.contrib.auth.backends import ModelBackend
 from ninja import Schema
-from ninja_jwt.schema import TokenObtainInputSchemaBase
+from ninja_jwt.schema import (
+    TokenObtainInputSchemaBase,
+    TokenRefreshInputSchema,
+    TokenRefreshOutputSchema,
+)
 from ninja_jwt.tokens import RefreshToken
 
 from core.schemas import StandResponse
@@ -22,6 +26,25 @@ class MyTokenObtainPairOutSchema(Schema):
     user: UserSchema
 
 
+class MyRefreshTokenOutSchema(TokenRefreshOutputSchema):
+    def dict(
+        self,
+        *,
+        include: Optional[Union["AbstractSetIntStr", "MappingIntStrAny"]] = None,
+        exclude: Optional[Union["AbstractSetIntStr", "MappingIntStrAny"]] = None,
+        by_alias: bool = False,
+        skip_defaults: Optional[bool] = None,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = False,
+    ) -> "DictStrAny":
+        return {
+            "data": {"access": self.access, "refresh": self.refresh},
+            "message": None,
+            "success": True,
+        }
+
+
 class MyTokenObtainPairInputSchema(TokenObtainInputSchemaBase):
     @classmethod
     def get_response_schema(cls) -> Type[Schema]:
@@ -35,10 +58,14 @@ def get_token(cls, user) -> Dict:
         refresh = RefreshToken.for_user(user)
         values["refresh"] = str(refresh)
         values["access"] = str(refresh.access_token)
-        values.update(
-            user=UserSchema.from_orm(user)
-        )
-        return {'data': values}
+        values.update(user=UserSchema.from_orm(user))
+        return {"data": values}
+
+
+class MyTokenRefreshInputSchema(TokenRefreshInputSchema):
+    @classmethod
+    def get_response_schema(cls) -> Type[Schema]:
+        return MyRefreshTokenOutSchema
 
 
 class CustomAuthBackend(ModelBackend):