Update main.py

main.py

@@ -12,8 +12,8 @@ from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 from datetime import datetime
 
 from aiohttp import ClientSession, ClientTimeout, ClientError
-from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi import FastAPI, HTTPException, Request, Depends, Header, Response
+from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse, PlainTextResponse
 from pydantic import BaseModel
 
 # Configure logging
@@ -43,9 +43,9 @@ else:
 rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 ip_rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 
-async def get_api_key(authorization: str = Header(...)) -> str:
-    if not authorization.startswith('Bearer '):
-        logger.warning("Invalid authorization header format.")
+async def get_api_key(authorization: str = Header(None)) -> str:
+    if authorization is None or not authorization.startswith('Bearer '):
+        logger.warning("Invalid or missing authorization header.")
         raise HTTPException(status_code=401, detail='Invalid authorization header format')
     api_key = authorization[7:]
     if api_key not in API_KEYS:
@@ -202,7 +202,7 @@ class Blackbox:
     }
 
     @classmethod
-    def get_model(cls, model: str) -> str:
+    def get_model(cls, model: str) -> Optional[str]:
         if model in cls.models:
             return model
         elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
@@ -373,6 +373,27 @@ class Blackbox:
 # FastAPI app setup
 app = FastAPI()
 
+# Middleware to enhance security
+@app.middleware("http")
+async def security_middleware(request: Request, call_next):
+    # Enforce that POST requests to sensitive endpoints must have a valid Content-Type
+    if request.method == "POST" and request.url.path in ["/v1/chat/completions", "/v1/completions"]:
+        content_type = request.headers.get("Content-Type")
+        if content_type != "application/json":
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": {
+                        "message": "Content-Type must be application/json",
+                        "type": "invalid_request_error",
+                        "param": None,
+                        "code": None
+                    }
+                },
+            )
+    response = await call_next(request)
+    return response
+
 class Message(BaseModel):
     role: str
     content: str
@@ -492,6 +513,12 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         logger.exception("An unexpected error occurred while processing the chat completions request.")
         raise HTTPException(status_code=500, detail=str(e))
 
+# Return 'about:blank' when accessing the endpoint via GET
+@app.get("/v1/chat/completions")
+async def chat_completions_get():
+    logger.info("GET request made to /v1/chat/completions, returning 'about:blank'")
+    return HTMLResponse(content='', status_code=200, media_type='text/html')
+
 @app.get("/v1/models")
 async def get_models():
     logger.info("Fetching available models")