new-auth

auth.py

@@ -0,0 +1,12 @@
+from typing import Optional
+
+def isAuthorized(authorization: Optional[str], TOKEN: str) -> bool: 
+    if authorization is None:
+        return False
+
+    if not authorization.startswith("Token "):
+        return False
+    
+    token = authorization.split(" ")[1]
+
+    return token == TOKEN

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, Header
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
@@ -10,6 +10,14 @@ from pydantic import BaseModel
 from typing import Optional
 from fastapi import FastAPI, HTTPException
 import json
+import os
+from auth import isAuthorized
+
+PASSCODE = os.getenv("PASSCODE")
+TOKEN = os.getenv("TOKEN")
+
+if PASSCODE == None or TOKEN == None:
+    raise Exception("Please set the PASSCODE and TOKEN environment variables.")
 
 app = FastAPI(
     title="DOCUMANTICAI API",
@@ -20,11 +28,14 @@ app = FastAPI(
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins="*"
+    allow_origins=["*"],
+    allow_headers=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
 )
 
 @app.post("/upload")
-async def upload_image(fields:str, model:str, file: UploadFile = File(...)):
+async def upload_image(fields:str, model:str, file: UploadFile = File(...), authorization: Optional[str] = Header(None)):
     """
     ### Endpoint Description:
     Extract form data from an uploaded image and return the extracted data in JSON format.
@@ -39,6 +50,9 @@ async def upload_image(fields:str, model:str, file: UploadFile = File(...)):
     """
     
     try:
+        if not isAuthorized(authorization, TOKEN):
+            return JSONResponse(content="Unauthorized", status_code=401)
+
         # Load the uploaded image
         image = Image.open(io.BytesIO(await file.read()))
         
@@ -121,7 +135,7 @@ class TranslateRequest(BaseModel):
     target_language: str
 
 @app.post("/translate")
-async def translate_text_endpoint(request: TranslateRequest):
+async def translate_text_endpoint(request: TranslateRequest, authorization: Optional[str] = Header(None)):
     """
     ### Endpoint Description:
     Translate text to a specified language.
@@ -134,6 +148,9 @@ async def translate_text_endpoint(request: TranslateRequest):
     - The translated text in the specified language.
     """
     try:
+        if not isAuthorized(authorization, TOKEN):
+            return JSONResponse(content="Unauthorized", status_code=401)
+
         response = translate_text(request.text, request.target_language)
         response = json.loads(response)
         return JSONResponse(content=response)
@@ -142,7 +159,7 @@ async def translate_text_endpoint(request: TranslateRequest):
 
 
 # @app.get("/translate")
-async def translate_text_endpoint(text: str, target_language: str):
+async def translate_text_endpoint(text: str, target_language: str, authorization: Optional[str] = Header(None)):
     """
     ### Endpoint Description:
     Translate text to a specified language.
@@ -156,6 +173,9 @@ async def translate_text_endpoint(text: str, target_language: str):
     """
 
     try:
+        if not isAuthorized(authorization, TOKEN):
+            return JSONResponse(content="Unauthorized", status_code=401)
+
         # print(target_language)
         response = translate_text(text, target_language)
         # print(response) #{"translated_text":"नमस्ते","translated_language":"नेपाली"}
@@ -170,11 +190,22 @@ class ChatRequest(BaseModel):
     history: Optional[str] = None
 
 @app.post("/chat")
-def chat_endpoint(request: ChatRequest):
+def chat_endpoint(request: ChatRequest, authorization: Optional[str] = Header(None)):
+    if not isAuthorized(authorization, TOKEN):
+        return JSONResponse(content="Unauthorized", status_code=401)
+
     response_text = chat_text(request.text, request.query, request.history)
     return {"reply": response_text}
 
 # @app.get("/chat")
 def chat_endpoint(query: str, text: str, history: Optional[str] = None):
     response_text = chat_text(text, query, history)
-    return {"reply": response_text}
+    return {"reply": response_text}
+
+@app.get('/token')
+def get_token(passcode: str):
+    if passcode == PASSCODE:
+       return {"token": TOKEN}
+    else:
+        return JSONResponse(content="invalid passcode", status_code=401)
+