From 360e6ac4ba1203527c484c59d37365b0ec47d70c Mon Sep 17 00:00:00 2001
From: Iam54r1n4 <iam54r1n40@proton.me>
Date: Tue, 4 Feb 2025 16:34:52 +0000
Subject: [PATCH] Implement custom url_for to generating urls with ROOT_PATH
 prefix

---
 core/scripts/webpanel/dependency/__init__.py  |  2 +-
 .../scripts/webpanel/dependency/dependency.py | 19 +++++++++++++++++++
 core/scripts/webpanel/middleware/auth.py      |  8 +++++---
 core/scripts/webpanel/routers/login/login.py  |  8 ++++----
 4 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/core/scripts/webpanel/dependency/__init__.py b/core/scripts/webpanel/dependency/__init__.py
index c4881d5..103eef4 100644
--- a/core/scripts/webpanel/dependency/__init__.py
+++ b/core/scripts/webpanel/dependency/__init__.py
@@ -1 +1 @@
-from .dependency import get_templates, get_session_manager
+from .dependency import get_templates, get_session_manager, url_for
diff --git a/core/scripts/webpanel/dependency/dependency.py b/core/scripts/webpanel/dependency/dependency.py
index 4f3d92c..b2feac3 100644
--- a/core/scripts/webpanel/dependency/dependency.py
+++ b/core/scripts/webpanel/dependency/dependency.py
@@ -1,4 +1,8 @@
+from fastapi import Request
 from fastapi.templating import Jinja2Templates
+from jinja2 import pass_context
+from typing import Any
+from starlette.datastructures import URL
 
 from session import SessionStorage, SessionManager
 from config import CONFIGS
@@ -6,6 +10,21 @@ from config import CONFIGS
 __TEMPLATES = Jinja2Templates(directory='templates')
 
 
+@pass_context
+def url_for(context: dict[str, Any], name: str = '', **path_params: dict[str, Any]) -> URL:
+    '''
+    Custom url_for function for Jinja2 to add a prefix to the generated URL.
+    '''
+    request: Request = context["request"]
+    url = request.url_for(name, **path_params)
+    prefixed_path = f"{CONFIGS.ROOT_PATH.rstrip('/')}/{url.path.lstrip('/')}"
+
+    return url.replace(path=prefixed_path)
+
+
+__TEMPLATES.env.globals['url_for'] = url_for  # type: ignore
+
+
 def get_templates() -> Jinja2Templates:
     return __TEMPLATES
 
diff --git a/core/scripts/webpanel/middleware/auth.py b/core/scripts/webpanel/middleware/auth.py
index 18ea388..78ab81f 100644
--- a/core/scripts/webpanel/middleware/auth.py
+++ b/core/scripts/webpanel/middleware/auth.py
@@ -5,6 +5,7 @@ from starlette.types import ASGIApp
 from typing import Awaitable, Callable
 from datetime import datetime, timezone
 
+from dependency import url_for
 from session import SessionManager
 
 
@@ -15,6 +16,7 @@ class AuthMiddleware(BaseHTTPMiddleware):
         super().__init__(app)
         self.__session_manager = session_manager
         self.__api_token = api_token
+        self.__url_for = url_for
 
     async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]):
         '''Handles session authentication.'''
@@ -41,7 +43,7 @@ class AuthMiddleware(BaseHTTPMiddleware):
         if not session_id:
             if is_api_request:
                 raise HTTPException(status_code=401, detail="Unauthorized")
-            return RedirectResponse(url='/login', status_code=302)
+            return RedirectResponse(url=self.__url_for(context={'request': request}, name='login'), status_code=302)
 
         session_data = self.__session_manager.get_session(session_id)
 
@@ -49,12 +51,12 @@ class AuthMiddleware(BaseHTTPMiddleware):
             if is_api_request:
                 raise HTTPException(status_code=401, detail="The session is invalid.")
 
-            return RedirectResponse(url='/login', status_code=302)
+            return RedirectResponse(url=self.__url_for(context={'request': request}, name='login'), status_code=302)
 
         if session_data.expires_at < datetime.now(timezone.utc):
             if is_api_request:
                 raise HTTPException(status_code=401, detail="The session has expired.")
 
-            return RedirectResponse(url='/login', status_code=302)
+            return RedirectResponse(url=self.__url_for(context={'request': request}, name='login'), status_code=302)
 
         return await call_next(request)
diff --git a/core/scripts/webpanel/routers/login/login.py b/core/scripts/webpanel/routers/login/login.py
index cdedf5c..5edd7b0 100644
--- a/core/scripts/webpanel/routers/login/login.py
+++ b/core/scripts/webpanel/routers/login/login.py
@@ -2,7 +2,7 @@ from fastapi import APIRouter, Depends, Form, Request
 from fastapi.responses import RedirectResponse
 from fastapi.templating import Jinja2Templates
 
-from dependency import get_templates, get_session_manager
+from dependency import get_templates, get_session_manager, url_for
 from session import SessionManager
 from config import CONFIGS
 
@@ -18,7 +18,7 @@ async def login(request: Request, templates: Jinja2Templates = Depends(get_templ
 async def login_post(
     request: Request,
     templates: Jinja2Templates = Depends(get_templates), session_manager: SessionManager = Depends(get_session_manager),
-    username: str = Form(), password: str = Form(),
+    username: str = Form(), password: str = Form()
 ):
     ADMIN_USERNAME = CONFIGS.ADMIN_USERNAME
     ADMIN_PASSWORD = CONFIGS.ADMIN_PASSWORD
@@ -28,7 +28,7 @@ async def login_post(
 
     session_id = session_manager.set_session(username)
 
-    res = RedirectResponse(url='/', status_code=302)
+    res = RedirectResponse(url=url_for(context={'request': request}, name='index'), status_code=302)
     res.set_cookie(key='session_id', value=session_id)
 
     return res
@@ -40,6 +40,6 @@ async def logout(request: Request, session_manager: SessionManager = Depends(get
     if session_id:
         session_manager.revoke_session(session_id)
 
-    res = RedirectResponse(url='/', status_code=302)
+    res = RedirectResponse(url=url_for(context={'request': request}, name='index'), status_code=302)
     res.delete_cookie('session_id')
     return res