autehntication

src/server/main.py

@@ -8,7 +8,7 @@ import uvicorn
 from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Form, Security
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
-from fastapi.security import OAuth2PasswordBearer
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, field_validator, Field
 from slowapi import Limiter
 from slowapi.util import get_remote_address
@@ -16,7 +16,7 @@ import requests
 from PIL import Image
 
 # Import from auth.py
-from utils.auth import get_current_user, login, refresh_token, register, TokenResponse, Settings, LoginRequest, RegisterRequest
+from utils.auth import get_current_user, get_current_user_with_admin, login, refresh_token, register, TokenResponse, Settings, LoginRequest, RegisterRequest, bearer_scheme
 
 # Assuming these are in your project structure
 from config.tts_config import SPEED, ResponseFormat, config as tts_config
@@ -24,18 +24,15 @@ from config.logging_config import logger
 
 settings = Settings()
 
-# Define OAuth2 scheme explicitly for Swagger
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/v1/token")
-
 # FastAPI app setup with enhanced docs
 app = FastAPI(
     title="Dhwani API",
     description="A multilingual AI-powered API supporting Indian languages for chat, text-to-speech, audio processing, and transcription. "
                 "**Authentication Guide:** \n"
-                "1. Register a new user via `/v1/register` with a POST request containing `username` and `password`. \n"
+                "1. Register a new user via `/v1/register` with a POST request containing `username` and `password` (requires admin access). \n"
                 "2. Obtain an access token by sending a POST request to `/v1/token` with `username` and `password`. \n"
-                "3. Click the 'Authorize' button (top-right), enter `Bearer <your_access_token>` in the 'bearerAuth' field, and click 'Authorize'. \n"
-                "All protected endpoints require this token for access.",
+                "3. Click the 'Authorize' button (top-right), enter your access token (e.g., `your_access_token`) in the 'bearerAuth' field, and click 'Authorize'. \n"
+                "All protected endpoints require this token for access. Only the 'admin' user (default password: adminpass) can register new users.",
     version="1.0.0",
     redirect_slashes=False,
     openapi_tags=[
@@ -45,7 +42,6 @@ app = FastAPI(
         {"name": "Audio", "description": "Audio processing and TTS endpoints"},
         {"name": "Translation", "description": "Text translation endpoints"}
     ],
-    swagger_ui_oauth2_redirect_url="/docs/oauth2-redirect",
 )
 
 app.add_middleware(
@@ -170,20 +166,24 @@ async def token(login_request: LoginRequest):
               200: {"description": "New token issued", "model": TokenResponse},
               401: {"description": "Invalid or expired token"}
           })
-async def refresh(token_response: TokenResponse = Depends(refresh_token)):
-    return token_response
+async def refresh(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)):
+    return await refresh_token(credentials)
 
 @app.post("/v1/register", 
           response_model=TokenResponse,
           summary="Register New User",
-          description="Create a new user account and return an access token. Use this token for subsequent requests.",
+          description="Create a new user account and return an access token. Requires admin access (use 'admin' user with password 'adminpass' initially).",
           tags=["Authentication"],
           responses={
               200: {"description": "User registered successfully", "model": TokenResponse},
-              400: {"description": "Username already exists"}
+              400: {"description": "Username already exists"},
+              403: {"description": "Admin access required"}
           })
-async def register_user(register_request: RegisterRequest):
-    return await register(register_request)
+async def register_user(
+    register_request: RegisterRequest,
+    current_user: str = Depends(get_current_user_with_admin)  # Enforce admin-only access
+):
+    return await register(register_request, current_user)  # Pass current_user explicitly
 
 @app.post("/v1/audio/speech",
           summary="Generate Speech from Text",
@@ -200,9 +200,10 @@ async def register_user(register_request: RegisterRequest):
 async def generate_audio(
     request: Request,
     speech_request: SpeechRequest = Depends(),
-    user_id: str = Depends(get_current_user),
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme),
     tts_service: TTSService = Depends(get_tts_service)
 ):
+    user_id = await get_current_user(credentials)
     if not speech_request.input.strip():
         raise HTTPException(status_code=400, detail="Input cannot be empty")
     
@@ -276,8 +277,9 @@ class ChatResponse(BaseModel):
 async def chat(
     request: Request,
     chat_request: ChatRequest,
-    user_id: str = Depends(get_current_user)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
 ):
+    user_id = await get_current_user(credentials)
     if not chat_request.prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, user_id: {user_id}")
@@ -332,8 +334,9 @@ async def process_audio(
     request: Request,
     file: UploadFile = File(..., description="Audio file to process"),
     language: str = Query(..., enum=["kannada", "hindi", "tamil"], description="Language of the audio"),
-    user_id: str = Depends(get_current_user),
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
 ):
+    user_id = await get_current_user(credentials)
     logger.info("Processing audio processing request", extra={
         "endpoint": "/v1/process_audio",
         "filename": file.filename,
@@ -378,9 +381,9 @@ async def process_audio(
 async def transcribe_audio(
     file: UploadFile = File(..., description="Audio file to transcribe"),
     language: str = Query(..., enum=["kannada", "hindi", "tamil"], description="Language of the audio"),
-    user_id: str = Depends(get_current_user),
-    request: Request = None,
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
 ):
+    user_id = await get_current_user(credentials)
     start_time = time()
     try:
         file_content = await file.read()
@@ -419,8 +422,9 @@ async def chat_v2(
     request: Request,
     prompt: str = Form(..., description="Text prompt for chat"),
     image: UploadFile = File(default=None, description="Optional image to accompany the prompt"),
-    user_id: str = Depends(get_current_user)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
 ):
+    user_id = await get_current_user(credentials)
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     
@@ -473,8 +477,9 @@ class TranslationResponse(BaseModel):
           })
 async def translate(
     request: TranslationRequest,
-    user_id: str = Depends(get_current_user)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
 ):
+    user_id = await get_current_user(credentials)
     logger.info(f"Received translation request: {request.dict()}, user_id: {user_id}")
     
     external_url = f"https://slabstech-dhwani-internal-api-server.hf.space/translate?src_lang={request.src_lang}&tgt_lang={request.tgt_lang}"

src/server/utils/auth.py

@@ -1,11 +1,11 @@
 import jwt
 from datetime import datetime, timedelta
-from fastapi.security import OAuth2PasswordBearer
 from fastapi import HTTPException, status, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings
 from config.logging_config import logger
-from sqlalchemy import create_engine, Column, String
+from sqlalchemy import create_engine, Column, String, Boolean
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker
 from passlib.context import CryptContext
@@ -20,6 +20,7 @@ class User(Base):
     __tablename__ = "users"
     username = Column(String, primary_key=True, index=True)
     password = Column(String)  # Stores hashed passwords
+    is_admin = Column(Boolean, default=False)  # New admin flag
 
 Base.metadata.create_all(bind=engine)
 
@@ -31,7 +32,11 @@ def seed_initial_data():
     db = SessionLocal()
     if not db.query(User).filter_by(username="testuser").first():
         hashed_password = pwd_context.hash("password123")
-        db.add(User(username="testuser", password=hashed_password))
+        db.add(User(username="testuser", password=hashed_password, is_admin=False))
+        db.commit()
+    if not db.query(User).filter_by(username="admin").first():
+        hashed_password = pwd_context.hash("adminpass")
+        db.add(User(username="admin", password=hashed_password, is_admin=True))
         db.commit()
     db.close()
 
@@ -58,7 +63,8 @@ class Settings(BaseSettings):
 settings = Settings()
 logger.info(f"Loaded API_KEY_SECRET at startup: {settings.api_key_secret}")
 
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/v1/token")
+# Use HTTPBearer
+bearer_scheme = HTTPBearer()
 
 class TokenPayload(BaseModel):
     sub: str
@@ -84,7 +90,8 @@ async def create_access_token(user_id: str) -> str:
     logger.info(f"Generated access token for user: {user_id}")
     return token
 
-async def get_current_user(token: str = Depends(oauth2_scheme)) -> str:
+async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)) -> str:
+    token = credentials.credentials
     credentials_exception = HTTPException(
         status_code=status.HTTP_401_UNAUTHORIZED,
         detail="Invalid authentication credentials",
@@ -127,6 +134,16 @@ async def get_current_user(token: str = Depends(oauth2_scheme)) -> str:
         logger.error(f"Unexpected token validation error: {str(e)}")
         raise credentials_exception
 
+async def get_current_user_with_admin(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)) -> str:
+    user_id = await get_current_user(credentials)
+    db = SessionLocal()
+    user = db.query(User).filter_by(username=user_id).first()
+    db.close()
+    if not user or not user.is_admin:
+        logger.warning(f"User {user_id} is not authorized as admin")
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Admin access required")
+    return user_id
+
 async def login(login_request: LoginRequest) -> TokenResponse:
     db = SessionLocal()
     user = db.query(User).filter_by(username=login_request.username).first()
@@ -137,7 +154,7 @@ async def login(login_request: LoginRequest) -> TokenResponse:
     token = await create_access_token(user_id=user.username)
     return TokenResponse(access_token=token, token_type="bearer")
 
-async def register(register_request: RegisterRequest) -> TokenResponse:
+async def register(register_request: RegisterRequest, current_user: str = Depends(get_current_user_with_admin)) -> TokenResponse:
     db = SessionLocal()
     existing_user = db.query(User).filter_by(username=register_request.username).first()
     if existing_user:
@@ -146,16 +163,16 @@ async def register(register_request: RegisterRequest) -> TokenResponse:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Username already exists")
     
     hashed_password = pwd_context.hash(register_request.password)
-    new_user = User(username=register_request.username, password=hashed_password)
+    new_user = User(username=register_request.username, password=hashed_password, is_admin=False)
     db.add(new_user)
     db.commit()
     db.close()
     
     token = await create_access_token(user_id=register_request.username)
-    logger.info(f"Registered and generated token for user: {register_request.username}")
+    logger.info(f"Registered and generated token for user: {register_request.username} by admin {current_user}")
     return TokenResponse(access_token=token, token_type="bearer")
 
-async def refresh_token(token: str = Depends(oauth2_scheme)) -> TokenResponse:
-    user_id = await get_current_user(token)
+async def refresh_token(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)) -> TokenResponse:
+    user_id = await get_current_user(credentials)
     new_token = await create_access_token(user_id=user_id)
     return TokenResponse(access_token=new_token, token_type="bearer")