main.py

from fastapi import FastAPI, Response, WebSocket, WebSocketDisconnect, Depends, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from dotenv import load_dotenv
import os

# Load environment variables from .env file
load_dotenv()
IS_DEV = os.environ.get('ENV', 'DEV') != 'PROD'
SECURE_TOKEN = os.getenv("SECURE_TOKEN")
X_REQUEST_USER = os.environ.get('X_REQUEST_USER')
X_API_KEY = os.environ.get('X_API_KEY')

app = FastAPI()

# CORS Middleware: restrict access to only trusted origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    #allow_origins=["https://your-frontend-domain.com"],
    #allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Simple token-based authentication dependency
async def authenticate_token(token: str):
    if token != f"Bearer {SECURE_TOKEN}":
        raise HTTPException(status_code=401, detail="Invalid token")

@app.get("/")
def root():
    return Response(status=200, data='ok')
        
@app.get("/health")
def healthcheck():
    return Response(status=200, data='ok')
        
# WebSocket endpoint
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = Depends(authenticate_token)):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            print(f"Message from client: {data}")
            await websocket.send_text(f"Server received: {data}")
    except WebSocketDisconnect:
        print("Client disconnected")

def is_valid(u, p):
    return u == X_REQUEST_USER and p == X_API_KEY

if __name__ == "__main__":
    uvicorn.run('app:app', host='0.0.0.0', port=7860, reload=True)