Update app.py

app.py

@@ -1,5 +1,6 @@
 import asyncio
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
+import os
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
@@ -11,13 +12,12 @@ import uvicorn
 app = FastAPI()
 
 # Configurer CORS pour autoriser toutes les origines
-# Permet à n'importe quel site web d'appeler votre API
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Autorise toutes les origines
+    allow_origins=["*"],
     allow_credentials=True,
-    allow_methods=["*"],  # Autorise toutes les méthodes (GET, POST, etc.)
-    allow_headers=["*"],  # Autorise tous les en-têtes
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
 # Monter un répertoire statique pour servir le fichier index.html
@@ -26,6 +26,7 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 class MockRequest(BaseModel):
     """Définit la structure attendue pour le corps de la requête POST."""
     parameter: str
+    secret: str  # Nouveau champ pour le secret
 
 class ConnectionManager:
     """Gère les connexions WebSocket actives."""
@@ -53,7 +54,7 @@ class ConnectionManager:
             await websocket.send_text(message)
             # Créer un Future pour attendre la réponse
             future = asyncio.get_event_loop().create_future()
-            # Utilise l'identifiant du client comme clé, bien que simple, c'est plus robuste
+            # Utilise l'identifiant du client comme clé
             client_id = str(id(websocket))
             self.response_futures[client_id] = future
             return future
@@ -61,6 +62,15 @@ class ConnectionManager:
 
 manager = ConnectionManager()
 
+def verify_secret(provided_secret: str) -> bool:
+    """Vérifie si le secret fourni correspond à celui de la variable d'environnement."""
+    expected_secret = os.getenv("API_SECRET")
+    
+    if not expected_secret:
+        print("ATTENTION: Variable d'environnement API_SECRET non définie!")
+        return False
+    
+    return provided_secret == expected_secret
 
 @app.get("/", response_class=HTMLResponse)
 async def root():
@@ -71,32 +81,38 @@ async def root():
     except FileNotFoundError:
         raise HTTPException(status_code=404, detail="index.html not found")
 
-
 @app.post("/v1/mock")
 async def mock_endpoint(payload: MockRequest):
     """
-    Endpoint API qui prend un string, le transmet via WebSocket,
-    attend une réponse et la retourne.
+    Endpoint API qui prend un string et un secret, vérifie le secret,
+    puis transmet via WebSocket, attend une réponse et la retourne.
     """
     try:
-        # Récupérer les données JSON du corps de la requête
-        # data = await request.json()
         input_string = payload.parameter
+        provided_secret = payload.secret
+
+        # Vérification du secret AVANT tout traitement
+        if not verify_secret(provided_secret):
+            print(f"Tentative d'accès avec un secret invalide: '{provided_secret[:10]}...'")
+            raise HTTPException(
+                status_code=401, 
+                detail="Secret invalide. Accès non autorisé."
+            )
+
+        print(f"Secret vérifié avec succès. Endpoint /v1/mock appelé avec: '{input_string}'")
 
         if input_string is None:
-            return {"error": "Le paramètre 'parameter' est manquant."}
+            raise HTTPException(status_code=400, detail="Le paramètre 'parameter' est manquant.")
             
-        print(f"Endpoint /v1/mock appelé avec: '{input_string}'")
-
         if not manager.active_connections:
-            return {"error": "Aucun client WebSocket n'est connecté."}
+            raise HTTPException(status_code=503, detail="Aucun client WebSocket n'est connecté.")
 
         # Envoyer le message via WebSocket et obtenir un "future" pour la réponse
         print("Envoi du message au client WebSocket...")
         response_future = await manager.broadcast(input_string)
 
         if response_future is None:
-             return {"error": "Échec de la diffusion du message."}
+            raise HTTPException(status_code=500, detail="Échec de la diffusion du message.")
 
         try:
             # Attendre la réponse du client WebSocket avec un timeout de 60 secondes
@@ -106,12 +122,14 @@ async def mock_endpoint(payload: MockRequest):
 
         except asyncio.TimeoutError:
             print("Timeout: Aucune réponse du client WebSocket.")
-            return {"error": "Timeout: Le client n'a pas répondu à temps."}
+            raise HTTPException(status_code=408, detail="Timeout: Le client n'a pas répondu à temps.")
 
+    except HTTPException:
+        # Re-lever les HTTPException sans les wrapper
+        raise
     except Exception as e:
         print(f"Erreur dans /v1/mock: {e}")
-        return {"error": f"Une erreur interne est survenue: {str(e)}"}
-
+        raise HTTPException(status_code=500, detail=f"Une erreur interne est survenue: {str(e)}")
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
@@ -136,6 +154,5 @@ async def websocket_endpoint(websocket: WebSocket):
         print(f"Erreur dans le WebSocket: {e}")
         manager.disconnect(websocket)
 
-
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)