add-uits

src/server/main.py

@@ -20,6 +20,9 @@ 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
 
+# Import decryption utility
+from utils.crypto import decrypt_data
+
 # Assuming these are in your project structure
 from config.tts_config import SPEED, ResponseFormat, config as tts_config
 from config.logging_config import logger
@@ -114,6 +117,77 @@ class AudioProcessingResponse(BaseModel):
     class Config:
         schema_extra = {"example": {"result": "Processed audio output"}} 
 
+class ChatRequest(BaseModel):
+    prompt: str = Field(..., description="Base64-encoded encrypted prompt (max 1000 characters after decryption)")
+    src_lang: str = Field("kan_Knda", description="Source language code (default: Kannada)")
+
+    @field_validator("prompt")
+    def prompt_must_be_valid(cls, v):
+        try:
+            base64.b64decode(v)
+        except Exception:
+            raise ValueError("Prompt must be valid base64-encoded data")
+        return v
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "prompt": "base64_encoded_encrypted_prompt",
+                "src_lang": "kan_Knda"
+            }
+        }
+
+class ChatResponse(BaseModel):
+    response: str = Field(..., description="Generated chat response")
+
+    class Config:
+        schema_extra = {"example": {"response": "Hi there, I'm doing great!"}} 
+
+class TranslationRequest(BaseModel):
+    sentences: List[str] = Field(..., description="List of base64-encoded encrypted sentences")
+    src_lang: str = Field(..., description="Source language code")
+    tgt_lang: str = Field(..., description="Target language code")
+
+    @field_validator("sentences")
+    def sentences_must_be_valid(cls, v):
+        for sentence in v:
+            try:
+                base64.b64decode(sentence)
+            except Exception:
+                raise ValueError("Each sentence must be valid base64-encoded data")
+        return v
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "sentences": ["base64_encoded_encrypted_hello", "base64_encoded_encrypted_how_are_you"],
+                "src_lang": "en",
+                "tgt_lang": "kan_Knda"
+            }
+        }
+
+class TranslationResponse(BaseModel):
+    translations: List[str] = Field(..., description="Translated sentences")
+
+    class Config:
+        schema_extra = {"example": {"translations": ["ನಮಸ್ಕಾರ", "ನೀವು ಹೇಗಿದ್ದೀರಿ?"]}} 
+
+class VisualQueryRequest(BaseModel):
+    query: str
+    src_lang: str = "kan_Knda"
+    tgt_lang: str = "kan_Knda"
+
+    @field_validator("query")
+    def query_must_be_valid(cls, v):
+        try:
+            base64.b64decode(v)
+        except Exception:
+            raise ValueError("Query must be valid base64-encoded data")
+        return v
+
+class VisualQueryResponse(BaseModel):
+    answer: str
+
 # TTS Service Interface
 class TTSService(ABC):
     @abstractmethod
@@ -271,38 +345,14 @@ async def generate_audio(
         headers=headers
     )
 
-class ChatRequest(BaseModel):
-    prompt: str = Field(..., description="Text prompt for chat (max 1000 characters)")
-    src_lang: str = Field("kan_Knda", description="Source language code (default: Kannada)")
-
-    @field_validator("prompt")
-    def prompt_must_be_valid(cls, v):
-        if len(v) > 1000:
-            raise ValueError("Prompt cannot exceed 1000 characters")
-        return v.strip()
-
-    class Config:
-        schema_extra = {
-            "example": {
-                "prompt": "Hello, how are you?",
-                "src_lang": "kan_Knda"
-            }
-        }
-
-class ChatResponse(BaseModel):
-    response: str = Field(..., description="Generated chat response")
-
-    class Config:
-        schema_extra = {"example": {"response": "Hi there, I'm doing great!"}} 
-
 @app.post("/v1/chat", 
           response_model=ChatResponse,
           summary="Chat with AI",
-          description="Generate a chat response from a prompt in the specified language. Rate limited to 100 requests per minute per user. Requires authentication.",
+          description="Generate a chat response from an encrypted prompt in the specified language. Rate limited to 100 requests per minute per user. Requires authentication and X-Session-Key header.",
           tags=["Chat"],
           responses={
               200: {"description": "Chat response", "model": ChatResponse},
-              400: {"description": "Invalid prompt"},
+              400: {"description": "Invalid prompt or encrypted data"},
               401: {"description": "Unauthorized - Token required"},
               429: {"description": "Rate limit exceeded"},
               504: {"description": "Chat service timeout"}
@@ -311,17 +361,31 @@ class ChatResponse(BaseModel):
 async def chat(
     request: Request,
     chat_request: ChatRequest,
-    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme),
+    x_session_key: str = Header(..., alias="X-Session-Key")
 ):
     user_id = await get_current_user(credentials)
-    if not chat_request.prompt:
+    session_key = base64.b64decode(x_session_key)
+    
+    # Decrypt the prompt
+    try:
+        encrypted_prompt = base64.b64decode(chat_request.prompt)
+        decrypted_prompt = decrypt_data(encrypted_prompt, session_key).decode("utf-8")
+    except Exception as e:
+        logger.error(f"Prompt decryption failed: {str(e)}")
+        raise HTTPException(status_code=400, detail="Invalid encrypted prompt")
+    
+    if not decrypted_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}")
+    if len(decrypted_prompt) > 1000:
+        raise HTTPException(status_code=400, detail="Decrypted prompt cannot exceed 1000 characters")
+    
+    logger.info(f"Received prompt: {decrypted_prompt}, src_lang: {chat_request.src_lang}, user_id: {user_id}")
     
     try:
         external_url = "https://slabstech-dhwani-internal-api-server.hf.space/v1/chat"
         payload = {
-            "prompt": chat_request.prompt,
+            "prompt": decrypted_prompt,
             "src_lang": chat_request.src_lang,
             "tgt_lang": chat_request.src_lang
         }
@@ -405,22 +469,27 @@ async def process_audio(
 @app.post("/v1/transcribe/", 
           response_model=TranscriptionResponse,
           summary="Transcribe Audio File",
-          description="Transcribe an uploaded audio file into text in the specified language. Requires authentication.",
+          description="Transcribe an encrypted audio file into text in the specified language. Requires authentication and X-Session-Key header.",
           tags=["Audio"],
           responses={
               200: {"description": "Transcription result", "model": TranscriptionResponse},
+              400: {"description": "Invalid encrypted audio"},
               401: {"description": "Unauthorized - Token required"},
               504: {"description": "Transcription service timeout"}
           })
 async def transcribe_audio(
-    file: UploadFile = File(..., description="Audio file to transcribe"),
+    file: UploadFile = File(..., description="Encrypted audio file to transcribe"),
     language: str = Query(..., enum=["kannada", "hindi", "tamil"], description="Language of the audio"),
-    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme),
+    x_session_key: str = Header(..., alias="X-Session-Key")
 ):
     user_id = await get_current_user(credentials)
+    session_key = base64.b64decode(x_session_key)
+    
     start_time = time()
     try:
-        file_content = await file.read()
+        encrypted_content = await file.read()
+        file_content = decrypt_data(encrypted_content, session_key)
         files = {"file": (file.filename, file_content, file.content_type)}
         
         external_url = f"{settings.external_asr_url}/transcribe/?language={language}"
@@ -436,7 +505,10 @@ async def transcribe_audio(
         logger.info(f"Transcription completed in {time() - start_time:.2f} seconds")
         return TranscriptionResponse(text=transcription)
     
+    except HTTPException:
+        raise
     except requests.Timeout:
+        logger.error("Transcription service timed out")
         raise HTTPException(status_code=504, detail="Transcription service timeout")
     except requests.RequestException as e:
         logger.error(f"Transcription request failed: {str(e)}")
@@ -480,48 +552,43 @@ async def chat_v2(
         logger.error(f"Chat_v2 processing failed: {str(e)}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
-class TranslationRequest(BaseModel):
-    sentences: List[str] = Field(..., description="List of sentences to translate")
-    src_lang: str = Field(..., description="Source language code")
-    tgt_lang: str = Field(..., description="Target language code")
-
-    class Config:
-        schema_extra = {
-            "example": {
-                "sentences": ["Hello", "How are you?"],
-                "src_lang": "en",
-                "tgt_lang": "kan_Knda"
-            }
-        }
-
-class TranslationResponse(BaseModel):
-    translations: List[str] = Field(..., description="Translated sentences")
-
-    class Config:
-        schema_extra = {"example": {"translations": ["ನಮಸ್ಕಾರ", "ನೀವು ಹೇಗಿದ್ದೀರಿ?"]}} 
-
 @app.post("/v1/translate", 
           response_model=TranslationResponse,
           summary="Translate Text",
-          description="Translate a list of sentences from source to target language. Requires authentication.",
+          description="Translate a list of base64-encoded encrypted sentences from source to target language. Requires authentication and X-Session-Key header.",
           tags=["Translation"],
           responses={
               200: {"description": "Translation result", "model": TranslationResponse},
+              400: {"description": "Invalid encrypted sentences"},
               401: {"description": "Unauthorized - Token required"},
               500: {"description": "Translation service error"},
               504: {"description": "Translation service timeout"}
           })
 async def translate(
     request: TranslationRequest,
-    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme),
+    x_session_key: str = Header(..., alias="X-Session-Key")
 ):
     user_id = await get_current_user(credentials)
-    logger.info(f"Received translation request: {request.dict()}, user_id: {user_id}")
+    session_key = base64.b64decode(x_session_key)
+    
+    # Decrypt sentences
+    decrypted_sentences = []
+    for sentence in request.sentences:
+        try:
+            encrypted_sentence = base64.b64decode(sentence)
+            decrypted_sentence = decrypt_data(encrypted_sentence, session_key).decode("utf-8")
+            decrypted_sentences.append(decrypted_sentence)
+        except Exception as e:
+            logger.error(f"Sentence decryption failed: {str(e)}")
+            raise HTTPException(status_code=400, detail="Invalid encrypted sentence")
+    
+    logger.info(f"Received translation request: {decrypted_sentences}, src_lang: {request.src_lang}, tgt_lang: {request.tgt_lang}, 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}"
     
     payload = {
-        "sentences": request.sentences,
+        "sentences": decrypted_sentences,
         "src_lang": request.src_lang,
         "tgt_lang": request.tgt_lang
     }
@@ -541,7 +608,7 @@ async def translate(
         response_data = response.json()
         translations = response_data.get("translations", [])
         
-        if not translations or len(translations) != len(request.sentences):
+        if not translations or len(translations) != len(decrypted_sentences):
             logger.warning(f"Unexpected response format: {response_data}")
             raise HTTPException(status_code=500, detail="Invalid response from translation service")
         
@@ -558,28 +625,14 @@ async def translate(
         logger.error(f"Invalid JSON response: {str(e)}")
         raise HTTPException(status_code=500, detail="Invalid response format from translation service")
 
-class VisualQueryRequest(BaseModel):
-    query: str
-    src_lang: str = "kan_Knda"
-    tgt_lang: str = "kan_Knda"
-
-    @field_validator("query")
-    def query_must_be_valid(cls, v):
-        if len(v) > 1000:
-            raise ValueError("Query cannot exceed 1000 characters")
-        return v.strip()
-
-class VisualQueryResponse(BaseModel):
-    answer: str
-
 @app.post("/v1/visual_query", 
           response_model=VisualQueryResponse,
           summary="Visual Query with Image",
-          description="Process a visual query with an image and text question. Rate limited to 100 requests per minute per user. Requires authentication.",
+          description="Process a visual query with an encrypted text query and encrypted image. Rate limited to 100 requests per minute per user. Requires authentication and X-Session-Key header.",
           tags=["Chat"],
           responses={
               200: {"description": "Query response", "model": VisualQueryResponse},
-              400: {"description": "Invalid query"},
+              400: {"description": "Invalid query or encrypted data"},
               401: {"description": "Unauthorized - Token required"},
               429: {"description": "Rate limit exceeded"},
               504: {"description": "Visual query service timeout"}
@@ -587,19 +640,40 @@ class VisualQueryResponse(BaseModel):
 @limiter.limit(settings.chat_rate_limit)
 async def visual_query(
     request: Request,
-    query: str = Form(..., description="Text query for the visual content"),
-    file: UploadFile = File(..., description="Image file to analyze"),
+    query: str = Form(..., description="Base64-encoded encrypted text query"),
+    file: UploadFile = File(..., description="Encrypted image file to analyze"),
     src_lang: str = Query(default="kan_Knda", description="Source language code"),
     tgt_lang: str = Query(default="kan_Knda", description="Target language code"),
-    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)
+    credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme),
+    x_session_key: str = Header(..., alias="X-Session-Key")
 ):
     user_id = await get_current_user(credentials)
-    if not query.strip():
+    session_key = base64.b64decode(x_session_key)
+    
+    # Decrypt query
+    try:
+        encrypted_query = base64.b64decode(query)
+        decrypted_query = decrypt_data(encrypted_query, session_key).decode("utf-8")
+    except Exception as e:
+        logger.error(f"Query decryption failed: {str(e)}")
+        raise HTTPException(status_code=400, detail="Invalid encrypted query")
+    
+    if not decrypted_query.strip():
         raise HTTPException(status_code=400, detail="Query cannot be empty")
+    if len(decrypted_query) > 1000:
+        raise HTTPException(status_code=400, detail="Decrypted query cannot exceed 1000 characters")
+    
+    # Decrypt image
+    try:
+        encrypted_content = await file.read()
+        decrypted_content = decrypt_data(encrypted_content, session_key)
+    except Exception as e:
+        logger.error(f"Image decryption failed: {str(e)}")
+        raise HTTPException(status_code=400, detail="Invalid encrypted image")
     
     logger.info("Processing visual query request", extra={
         "endpoint": "/v1/visual_query",
-        "query_length": len(query),
+        "query_length": len(decrypted_query),
         "file_name": file.filename,
         "client_ip": get_remote_address(request),
         "user_id": user_id,
@@ -610,9 +684,8 @@ async def visual_query(
     external_url = f"https://slabstech-dhwani-internal-api-server.hf.space/v1/visual_query/?src_lang={src_lang}&tgt_lang={tgt_lang}"
     
     try:
-        file_content = await file.read()
-        files = {"file": (file.filename, file_content, file.content_type)}
-        data = {"query": query}
+        files = {"file": (file.filename, decrypted_content, file.content_type)}
+        data = {"query": decrypted_query}
         
         response = requests.post(
             external_url,
@@ -650,19 +723,9 @@ class SupportedLanguage(str, Enum):
     hindi = "hindi"
     tamil = "tamil"
 
-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 and X-Session-Key header.",
+          description="Convert input encrypted 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"}}},
@@ -693,7 +756,7 @@ async def speech_to_speech(
 
     try:
         encrypted_content = await file.read()
-        file_content = decrypt_audio(encrypted_content, session_key)
+        file_content = decrypt_data(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/crypto.py

@@ -0,0 +1,13 @@
+from Crypto.Cipher import AES
+from fastapi import HTTPException
+from config.logging_config import logger
+
+def decrypt_data(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"Decryption failed: {str(e)}")
+        raise HTTPException(status_code=400, detail="Invalid encrypted data")