add-encrity

src/server/main.py

@@ -5,7 +5,7 @@ from typing import List, Optional
 from abc import ABC, abstractmethod
 
 import uvicorn
-from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Form
+from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Form, Header
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
@@ -14,6 +14,8 @@ from slowapi import Limiter
 from slowapi.util import get_remote_address
 import requests
 from PIL import Image
+import base64
+from Crypto.Cipher import AES
 
 # Import from auth.py
 from utils.auth import get_current_user, get_current_user_with_admin, login, refresh_token, register, app_register, TokenResponse, Settings, LoginRequest, RegisterRequest, bearer_scheme
@@ -58,7 +60,7 @@ async def get_user_id_for_rate_limit(request: Request):
         user_id = await get_current_user(credentials)
         return user_id
     except Exception:
-        return get_remote_address(request)  # Fallback to IP if unauthenticated
+        return get_remote_address(request)
 
 limiter = Limiter(key_func=get_user_id_for_rate_limit)
 
@@ -155,14 +157,18 @@ async def home():
 @app.post("/v1/token", 
           response_model=TokenResponse,
           summary="User Login",
-          description="Authenticate a user with username and password to obtain an access token and refresh token. Copy the access token and use it in the 'Authorize' button above.",
+          description="Authenticate a user with encrypted email and device token to obtain an access token and refresh token. Requires X-Session-Key header.",
           tags=["Authentication"],
           responses={
               200: {"description": "Successful login", "model": TokenResponse},
-              401: {"description": "Invalid username or password"}
+              400: {"description": "Invalid encrypted data"},
+              401: {"description": "Invalid email or device token"}
           })
-async def token(login_request: LoginRequest):
-    return await login(login_request)
+async def token(
+    login_request: LoginRequest,
+    x_session_key: str = Header(..., alias="X-Session-Key")
+):
+    return await login(login_request, x_session_key)
 
 @app.post("/v1/refresh", 
           response_model=TokenResponse,
@@ -195,20 +201,21 @@ async def register_user(
 @app.post("/v1/app/register",
           response_model=TokenResponse,
           summary="Register New App User",
-          description="Create a new user account for the mobile app using an email and device token. Returns an access token and refresh token. Rate limited to 5 requests per minute per IP.",
+          description="Create a new user account for the mobile app using an encrypted email and device token. Returns an access token and refresh token. Rate limited to 5 requests per minute per IP. Requires X-Session-Key header.",
           tags=["Authentication"],
           responses={
               200: {"description": "User registered successfully", "model": TokenResponse},
-              400: {"description": "Email already registered"},
+              400: {"description": "Email already registered or invalid encrypted data"},
               429: {"description": "Rate limit exceeded"}
           })
 @limiter.limit(settings.speech_rate_limit)
 async def app_register_user(
     request: Request,
-    register_request: RegisterRequest
+    register_request: RegisterRequest,
+    x_session_key: str = Header(..., alias="X-Session-Key")
 ):
-    logger.info(f"App registration attempt for email: {register_request.username}")
-    return await app_register(register_request)
+    logger.info(f"App registration attempt")
+    return await app_register(register_request, x_session_key)
 
 @app.post("/v1/audio/speech",
           summary="Generate Speech from Text",
@@ -551,8 +558,8 @@ async def translate(
 
 class VisualQueryRequest(BaseModel):
     query: str
-    src_lang: str = "kan_Knda"  # Default to Kannada
-    tgt_lang: str = "kan_Knda"  # Default to Kannada
+    src_lang: str = "kan_Knda"
+    tgt_lang: str = "kan_Knda"
 
     @field_validator("query")
     def query_must_be_valid(cls, v):
@@ -633,25 +640,31 @@ async def visual_query(
     except ValueError as e:
         logger.error(f"Invalid JSON response: {str(e)}")
         raise HTTPException(status_code=500, detail="Invalid response format from visual query service")
-    
-# Ensure these imports are at the top with other imports
-from fastapi.responses import StreamingResponse
+
 from enum import Enum
 
-# Define supported languages for validation
 class SupportedLanguage(str, Enum):
     kannada = "kannada"
     hindi = "hindi"
     tamil = "tamil"
 
-# Add the speech-to-speech endpoint
+def decrypt_audio(encrypted_data: bytes, key: bytes) -> bytes:
+    try:
+        nonce, ciphertext = encrypted_data[:12], encrypted_data[12:]
+        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
+        plaintext = cipher.decrypt_and_verify(ciphertext[:-16], ciphertext[-16:])
+        return plaintext
+    except Exception as e:
+        logger.error(f"Audio decryption failed: {str(e)}")
+        raise HTTPException(status_code=400, detail="Invalid encrypted audio")
+
 @app.post("/v1/speech_to_speech",
           summary="Speech-to-Speech Conversion",
-          description="Convert input speech to processed speech by calling an external speech-to-speech API. Rate limited to 5 requests per minute per user. Requires authentication.",
+          description="Convert input speech to processed speech by calling an external speech-to-speech API. Rate limited to 5 requests per minute per user. Requires authentication and X-Session-Key header.",
           tags=["Audio"],
           responses={
               200: {"description": "Audio stream", "content": {"audio/mp3": {"example": "Binary audio data"}}},
-              400: {"description": "Invalid input"},
+              400: {"description": "Invalid input or encrypted audio"},
               401: {"description": "Unauthorized - Token required"},
               429: {"description": "Rate limit exceeded"},
               504: {"description": "External API timeout"},
@@ -660,11 +673,14 @@ class SupportedLanguage(str, Enum):
 @limiter.limit(settings.speech_rate_limit)
 async def speech_to_speech(
     request: Request,
-    file: UploadFile = File(..., description="Audio file to process"),
+    file: UploadFile = File(..., description="Encrypted audio file to process"),
     language: SupportedLanguage = Query(..., description="Language of the audio (kannada, hindi, tamil)"),
-    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme),
+    x_session_key: str = Header(..., alias="X-Session-Key")
 ) -> StreamingResponse:
     user_id = await get_current_user(credentials)
+    session_key = base64.b64decode(x_session_key)
+    
     logger.info("Processing speech-to-speech request", extra={
         "endpoint": "/v1/speech_to_speech",
         "audio_filename": file.filename,
@@ -674,7 +690,8 @@ async def speech_to_speech(
     })
 
     try:
-        file_content = await file.read()
+        encrypted_content = await file.read()
+        file_content = decrypt_audio(encrypted_content, session_key)
         files = {"file": (file.filename, file_content, file.content_type)}
         external_url = f"https://slabstech-dhwani-internal-api-server.hf.space/v1/speech_to_speech?language={language}"
 

src/server/utils/auth.py

@@ -10,6 +10,9 @@ from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker
 from passlib.context import CryptContext
 import os
+import base64
+from Crypto.Cipher import AES
+from Crypto.Random import get_random_bytes
 
 # SQLite database setup with Hugging Face persistent storage
 DATABASE_PATH = "/data/users.db"
@@ -22,7 +25,8 @@ 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
+    is_admin = Column(Boolean, default=False)
+    session_key = Column(String, nullable=True)  # Stores base64-encoded session key
 
 # Ensure the /data directory exists
 os.makedirs(os.path.dirname(DATABASE_PATH), exist_ok=True)
@@ -35,8 +39,8 @@ pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 class Settings(BaseSettings):
     api_key_secret: str = Field(..., env="API_KEY_SECRET")
-    token_expiration_minutes: int = Field(1440, env="TOKEN_EXPIRATION_MINUTES")  # 24 hours
-    refresh_token_expiration_days: int = Field(7, env="REFRESH_TOKEN_EXPIRATION_DAYS")  # 7 days
+    token_expiration_minutes: int = Field(1440, env="TOKEN_EXPIRATION_MINUTES")
+    refresh_token_expiration_days: int = Field(7, env="REFRESH_TOKEN_EXPIRATION_DAYS")
     llm_model_name: str = "google/gemma-3-4b-it"
     max_tokens: int = 512
     host: str = "0.0.0.0"
@@ -61,19 +65,19 @@ settings = Settings()
 def seed_initial_data():
     db = SessionLocal()
     try:
-        # Seed test user (non-admin) with a device token-like password
         test_username = "[email protected]"
         if not db.query(User).filter_by(username=test_username).first():
-            test_device_token = "550e8400-e29b-41d4-a716-446655440000"  # Sample UUID
+            test_device_token = "550e8400-e29b-41d4-a716-446655440000"
             hashed_password = pwd_context.hash(test_device_token)
-            db.add(User(username=test_username, password=hashed_password, is_admin=False))
+            session_key = base64.b64encode(get_random_bytes(16)).decode('utf-8')
+            db.add(User(username=test_username, password=hashed_password, is_admin=False, session_key=session_key))
             db.commit()
-        # Seed admin user using environment variables
         admin_username = settings.default_admin_username
         admin_password = settings.default_admin_password
         if not db.query(User).filter_by(username=admin_username).first():
             hashed_password = pwd_context.hash(admin_password)
-            db.add(User(username=admin_username, password=hashed_password, is_admin=True))
+            session_key = base64.b64encode(get_random_bytes(16)).decode('utf-8')
+            db.add(User(username=admin_username, password=hashed_password, is_admin=True, session_key=session_key))
             db.commit()
         logger.info(f"Seeded initial data: test user '{test_username}', admin user '{admin_username}'")
     except Exception as e:
@@ -82,10 +86,8 @@ def seed_initial_data():
     finally:
         db.close()
 
-# Initialize database with seed data
 seed_initial_data()
 
-# Use HTTPBearer
 bearer_scheme = HTTPBearer()
 
 class TokenPayload(BaseModel):
@@ -106,6 +108,17 @@ class RegisterRequest(BaseModel):
     username: str
     password: str
 
+def decrypt_data(encrypted_data: str, key: bytes) -> str:
+    try:
+        data = base64.b64decode(encrypted_data)
+        nonce, ciphertext = data[:12], data[12:]
+        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
+        plaintext = cipher.decrypt_and_verify(ciphertext[:-16], ciphertext[-16:])
+        return plaintext.decode('utf-8')
+    except Exception as e:
+        logger.error(f"Decryption failed: {str(e)}")
+        raise HTTPException(status_code=400, detail="Invalid encrypted data")
+
 async def create_access_token(user_id: str) -> dict:
     expire = datetime.utcnow() + timedelta(minutes=settings.token_expiration_minutes)
     payload = {"sub": user_id, "exp": expire.timestamp(), "type": "access"}
@@ -166,14 +179,28 @@ async def get_current_user_with_admin(credentials: HTTPAuthorizationCredentials
         raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Admin access required")
     return user_id
 
-async def login(login_request: LoginRequest) -> TokenResponse:
+async def login(login_request: LoginRequest, session_key_b64: str) -> TokenResponse:
     db = SessionLocal()
-    user = db.query(User).filter_by(username=login_request.username).first()
-    db.close()
-    if not user or not pwd_context.verify(login_request.password, user.password):
-        logger.warning(f"Login failed for user: {login_request.username}")
+    session_key = base64.b64decode(session_key_b64)
+    try:
+        username = decrypt_data(login_request.username, session_key)
+        password = decrypt_data(login_request.password, session_key)
+    except:
+        db.close()
+        raise HTTPException(status_code=400, detail="Invalid encrypted data")
+    
+    user = db.query(User).filter_by(username=username).first()
+    if not user or not pwd_context.verify(password, user.password):
+        db.close()
+        logger.warning(f"Login failed for user: {username}")
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid email or device token")
-    tokens = await create_access_token(user_id=user.username)
+    
+    if user.session_key != session_key_b64:
+        user.session_key = session_key_b64
+        db.commit()
+    db.close()
+    
+    tokens = await create_access_token(user_id=username)
     return TokenResponse(access_token=tokens["access_token"], refresh_token=tokens["refresh_token"], token_type="bearer")
 
 async def register(register_request: RegisterRequest, current_user: str = Depends(get_current_user_with_admin)) -> TokenResponse:
@@ -194,22 +221,30 @@ async def register(register_request: RegisterRequest, current_user: str = Depend
     logger.info(f"Registered and generated token for user: {register_request.username} by admin {current_user}")
     return TokenResponse(access_token=tokens["access_token"], refresh_token=tokens["refresh_token"], token_type="bearer")
 
-async def app_register(register_request: RegisterRequest) -> TokenResponse:
+async def app_register(register_request: RegisterRequest, session_key_b64: str) -> TokenResponse:
     db = SessionLocal()
-    existing_user = db.query(User).filter_by(username=register_request.username).first()
+    session_key = base64.b64decode(session_key_b64)
+    try:
+        username = decrypt_data(register_request.username, session_key)
+        password = decrypt_data(register_request.password, session_key)
+    except:
+        db.close()
+        raise HTTPException(status_code=400, detail="Invalid encrypted data")
+    
+    existing_user = db.query(User).filter_by(username=username).first()
     if existing_user:
         db.close()
-        logger.warning(f"App registration failed: Email {register_request.username} already exists")
+        logger.warning(f"App registration failed: Email {username} already exists")
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Email already registered")
     
-    hashed_password = pwd_context.hash(register_request.password)
-    new_user = User(username=register_request.username, password=hashed_password, is_admin=False)
+    hashed_password = pwd_context.hash(password)
+    new_user = User(username=username, password=hashed_password, is_admin=False, session_key=session_key_b64)
     db.add(new_user)
     db.commit()
     db.close()
     
-    tokens = await create_access_token(user_id=register_request.username)
-    logger.info(f"App registered new user: {register_request.username}")
+    tokens = await create_access_token(user_id=username)
+    logger.info(f"App registered new user: {username}")
     return TokenResponse(access_token=tokens["access_token"], refresh_token=tokens["refresh_token"], token_type="bearer")
 
 async def refresh_token(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)) -> TokenResponse: