From 440642651585129a20773e6409960252adbccefa Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Mon, 17 Mar 2025 11:51:30 +0530 Subject: [PATCH] feat: implement session_jwt and enhance auth --- application/app.py | 42 ++++++++++++++++-- application/auth.py | 23 +++------- application/core/settings.py | 1 - frontend/src/App.tsx | 43 +++++++++++++----- frontend/src/Navigation.tsx | 17 +++----- frontend/src/api/endpoints.ts | 2 + frontend/src/api/services/userService.ts | 3 ++ frontend/src/hooks/useTokenAuth.ts | 55 ++++++++++++++++++++++++ frontend/src/modals/JWTModal.tsx | 21 +++------ 9 files changed, 148 insertions(+), 59 deletions(-) create mode 100644 frontend/src/hooks/useTokenAuth.ts diff --git a/application/app.py b/application/app.py index ddac00b0..7ca0ac2b 100644 --- a/application/app.py +++ b/application/app.py @@ -1,10 +1,12 @@ +import os import platform +import uuid import dotenv from flask import Flask, jsonify, redirect, request from jose import jwt -from application.auth import get_or_create_user_id, handle_auth +from application.auth import handle_auth from application.core.logging_config import setup_logging @@ -38,10 +40,22 @@ app.config.update( celery.config_from_object("application.celeryconfig") api.init_app(app) +if settings.AUTH_TYPE in ("simple_jwt", "session_jwt") and not settings.JWT_SECRET_KEY: + key_file = ".jwt_secret_key" + try: + with open(key_file, "r") as f: + settings.JWT_SECRET_KEY = f.read().strip() + except FileNotFoundError: + new_key = os.urandom(32).hex() + with open(key_file, "w") as f: + f.write(new_key) + settings.JWT_SECRET_KEY = new_key + except Exception as e: + raise RuntimeError(f"Failed to setup JWT_SECRET_KEY: {e}") + SIMPLE_JWT_TOKEN = None if settings.AUTH_TYPE == "simple_jwt": - user_id = get_or_create_user_id() - payload = {"sub": user_id} + payload = {"sub": "local"} SIMPLE_JWT_TOKEN = jwt.encode(payload, settings.JWT_SECRET_KEY, algorithm="HS256") print(f"Generated Simple JWT Token: {SIMPLE_JWT_TOKEN}") @@ -54,13 +68,33 @@ def home(): return "Welcome to DocsGPT Backend!" +@app.route("/api/config") +def get_config(): + response = { + "auth_type": settings.AUTH_TYPE, + "requires_auth": settings.AUTH_TYPE in ["simple_jwt", "session_jwt"], + } + return jsonify(response) + + +@app.route("/api/generate_token") +def generate_token(): + if settings.AUTH_TYPE == "session_jwt": + new_user_id = str(uuid.uuid4()) + token = jwt.encode( + {"sub": new_user_id}, settings.JWT_SECRET_KEY, algorithm="HS256" + ) + return jsonify({"token": token}) + return jsonify({"error": "Token generation not allowed in current auth mode"}), 400 + + @app.before_request def authenticate_request(): if request.method == "OPTIONS": return "", 200 decoded_token = handle_auth(request) - if "message" in decoded_token: + if not decoded_token: request.decoded_token = None elif "error" in decoded_token: return jsonify(decoded_token), 401 diff --git a/application/auth.py b/application/auth.py index be2501b0..78926c45 100644 --- a/application/auth.py +++ b/application/auth.py @@ -1,15 +1,13 @@ -import uuid - from jose import jwt from application.core.settings import settings def handle_auth(request, data={}): - if settings.AUTH_TYPE == "simple_jwt": + if settings.AUTH_TYPE in ["simple_jwt", "session_jwt"]: jwt_token = request.headers.get("Authorization") if not jwt_token: - return {"message": "Missing Authorization header"} + return None jwt_token = jwt_token.replace("Bearer ", "") @@ -22,18 +20,9 @@ def handle_auth(request, data={}): ) return decoded_token except Exception as e: - return {"message": f"Authentication error: {str(e)}"} + return { + "message": f"Authentication error: {str(e)}", + "error": "invalid_token", + } else: return {"sub": "local"} - - -def get_or_create_user_id(): - try: - with open(settings.USER_ID_FILE, "r") as f: - user_id = f.read().strip() - return user_id - except FileNotFoundError: - user_id = str(uuid.uuid4()) - with open(settings.USER_ID_FILE, "w") as f: - f.write(user_id) - return user_id diff --git a/application/core/settings.py b/application/core/settings.py index a757a7f3..2a78b3a0 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -100,7 +100,6 @@ class Settings(BaseSettings): FLASK_DEBUG_MODE: bool = False JWT_SECRET_KEY: str = "" - USER_ID_FILE: str = os.path.join(current_dir, "user_id.txt") path = Path(__file__).parent.parent.absolute() diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index ba0a4bd7..64c4c486 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -1,15 +1,30 @@ -import { Routes, Route } from 'react-router-dom'; -import Navigation from './Navigation'; -import Conversation from './conversation/Conversation'; -import About from './About'; -import PageNotFound from './PageNotFound'; -import { useMediaQuery } from './hooks'; -import { useState } from 'react'; -import Setting from './settings'; import './locale/i18n'; -import { Outlet } from 'react-router-dom'; + +import { useState } from 'react'; +import { Outlet, Route, Routes } from 'react-router-dom'; + +import About from './About'; +import Spinner from './components/Spinner'; +import Conversation from './conversation/Conversation'; import { SharedConversation } from './conversation/SharedConversation'; -import { useDarkTheme } from './hooks'; +import { useDarkTheme, useMediaQuery } from './hooks'; +import useTokenAuth from './hooks/useTokenAuth'; +import Navigation from './Navigation'; +import PageNotFound from './PageNotFound'; +import Setting from './settings'; + +function AuthWrapper({ children }: { children: React.ReactNode }) { + const { isAuthLoading } = useTokenAuth(); + + if (isAuthLoading) { + return ( +
+ +
+ ); + } + return <>{children}; +} function MainLayout() { const { isMobile } = useMediaQuery(); @@ -39,7 +54,13 @@ export default function App() { return (
- }> + + + + } + > } /> } /> } /> diff --git a/frontend/src/Navigation.tsx b/frontend/src/Navigation.tsx index ad30228b..c0fb8231 100644 --- a/frontend/src/Navigation.tsx +++ b/frontend/src/Navigation.tsx @@ -28,6 +28,7 @@ import { import ConversationTile from './conversation/ConversationTile'; import { useDarkTheme, useMediaQuery } from './hooks'; import useDefaultDocument from './hooks/useDefaultDocument'; +import useTokenAuth from './hooks/useTokenAuth'; import DeleteConvModal from './modals/DeleteConvModal'; import JWTModal from './modals/JWTModal'; import { ActiveState, Doc } from './models/misc'; @@ -72,10 +73,10 @@ export default function Navigation({ navOpen, setNavOpen }: NavigationProps) { const { t } = useTranslation(); const isApiKeySet = useSelector(selectApiKeyStatus); + const { showTokenModal, handleTokenSubmit } = useTokenAuth(); + const [uploadModalState, setUploadModalState] = useState('INACTIVE'); - const [authKeyModalState, setAuthKeyModalState] = - useState('INACTIVE'); const navRef = useRef(null); @@ -204,15 +205,7 @@ export default function Navigation({ navOpen, setNavOpen }: NavigationProps) { setNavOpen(!isMobile); }, [isMobile]); - useEffect(() => { - const authToken = localStorage.getItem('authToken'); - if (!authToken) { - setAuthKeyModalState('ACTIVE'); - } - }, []); - useDefaultDocument(); - return ( <> {!navOpen && ( @@ -485,8 +478,8 @@ export default function Navigation({ navOpen, setNavOpen }: NavigationProps) { > )} ); diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 9bf659de..0d574f89 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -1,5 +1,7 @@ const endpoints = { USER: { + CONFIG: '/api/config', + NEW_TOKEN: '/api/generate_token', DOCS: '/api/sources', DOCS_CHECK: '/api/docs_check', DOCS_PAGINATED: '/api/sources/paginated', diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index fbaa06ba..13083677 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -2,6 +2,9 @@ import apiClient from '../client'; import endpoints from '../endpoints'; const userService = { + getConfig: (): Promise => apiClient.get(endpoints.USER.CONFIG, null), + getNewToken: (): Promise => + apiClient.get(endpoints.USER.NEW_TOKEN, null), getDocs: (token: string | null): Promise => apiClient.get(`${endpoints.USER.DOCS}`, token), getDocsWithPagination: (query: string, token: string | null): Promise => diff --git a/frontend/src/hooks/useTokenAuth.ts b/frontend/src/hooks/useTokenAuth.ts new file mode 100644 index 00000000..8f408600 --- /dev/null +++ b/frontend/src/hooks/useTokenAuth.ts @@ -0,0 +1,55 @@ +import { useEffect, useRef, useState } from 'react'; +import { useDispatch, useSelector } from 'react-redux'; + +import userService from '../api/services/userService'; +import { selectToken, setToken } from '../preferences/preferenceSlice'; + +export default function useAuth() { + const dispatch = useDispatch(); + const token = useSelector(selectToken); + const [authType, setAuthType] = useState(null); + const [showTokenModal, setShowTokenModal] = useState(false); + const [isAuthLoading, setIsAuthLoading] = useState(true); + const isGeneratingToken = useRef(false); + + const generateNewToken = async () => { + if (isGeneratingToken.current) return; + isGeneratingToken.current = true; + const response = await userService.getNewToken(); + const { token: newToken } = await response.json(); + localStorage.setItem('authToken', newToken); + dispatch(setToken(newToken)); + setIsAuthLoading(false); + return newToken; + }; + + useEffect(() => { + const initializeAuth = async () => { + try { + const configRes = await userService.getConfig(); + const config = await configRes.json(); + setAuthType(config.auth_type); + + if (config.auth_type === 'session_jwt' && !token) { + await generateNewToken(); + } else if (config.auth_type === 'simple_jwt' && !token) { + setShowTokenModal(true); + setIsAuthLoading(false); + } else { + setIsAuthLoading(false); + } + } catch (error) { + console.error('Auth initialization failed:', error); + setIsAuthLoading(false); + } + }; + initializeAuth(); + }, []); + + const handleTokenSubmit = (enteredToken: string) => { + localStorage.setItem('authToken', enteredToken); + dispatch(setToken(enteredToken)); + setShowTokenModal(false); + }; + return { authType, showTokenModal, isAuthLoading, token, handleTokenSubmit }; +} diff --git a/frontend/src/modals/JWTModal.tsx b/frontend/src/modals/JWTModal.tsx index 90a13254..1cab145e 100644 --- a/frontend/src/modals/JWTModal.tsx +++ b/frontend/src/modals/JWTModal.tsx @@ -3,30 +3,23 @@ import { useDispatch } from 'react-redux'; import Input from '../components/Input'; import { ActiveState } from '../models/misc'; -import { setToken } from '../preferences/preferenceSlice'; import WrapperModal from './WrapperModal'; type JWTModalProps = { modalState: ActiveState; - setModalState: (state: ActiveState) => void; + handleTokenSubmit: (enteredToken: string) => void; }; -export default function JWTModal({ modalState, setModalState }: JWTModalProps) { - const dispatch = useDispatch(); +export default function JWTModal({ + modalState, + handleTokenSubmit, +}: JWTModalProps) { const [jwtToken, setJwtToken] = useState(''); - const handleSaveToken = () => { - if (jwtToken) { - localStorage.setItem('authToken', jwtToken); - dispatch(setToken(jwtToken)); - setModalState('INACTIVE'); - } - }; - if (modalState !== 'ACTIVE') return null; return ( - setModalState('INACTIVE')} className="p-4"> + {}}>
Add JWT Token @@ -44,7 +37,7 @@ export default function JWTModal({ modalState, setModalState }: JWTModalProps) {