Update app.py

app.py

@@ -1,112 +1,172 @@
-from fastapi import FastAPI, HTTPException, Depends, Request, status, Form
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.security import OAuth2PasswordRequestForm
-from jose import JWTError, jwt
-from passlib.context import CryptContext
-from datetime import datetime, timedelta
-from pymongo import MongoClient
-from bson import ObjectId
-from fastapi.responses import JSONResponse
-from dotenv import load_dotenv
-import os
-
-load_dotenv()
-
-# Constants
-SECRET_KEY = os.environ.get("SECRET_KEY")
-ALGORITHM = os.environ.get("ALGORITHM")
-ACCESS_TOKEN_EXPIRE_MINUTES = 60
-
-# Admin credentials
-fake_admin_db = {
-    "admin": {
-        "username": os.environ.get("ADMIN_USERNAME"),
-        "hashed_password": os.environ.get("ADMIN_PASSWORD")
-    }
-}
-
-pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
-client = MongoClient(os.environ.get("MONGODB_URI"))
-db = client["thehexatechdb"]
-collection = db["quotationsdb"]
-
-app = FastAPI()
-
-# CORS
-origins = ["*"]
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Auth Utilities
-def verify_password(plain, hashed):
-    return pwd_context.verify(plain, hashed)
-
-def get_password_hash(password):
-    return pwd_context.hash(password)
-
-def authenticate_user(username: str, password: str):
-    user = fake_admin_db.get(username)
-    if not user or not verify_password(password, user["hashed_password"]):
-        return False
-    return {"username": username}
-
-def create_access_token(data: dict, expires_delta=None):
-    to_encode = data.copy()
-    expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
-    to_encode.update({"exp": expire})
-    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
-
-def get_current_user(token: str = Depends(OAuth2PasswordRequestForm)):
-    credentials_exception = HTTPException(
-        status_code=status.HTTP_401_UNAUTHORIZED,
-        detail="Invalid credentials",
-        headers={"WWW-Authenticate": "Bearer"},
-    )
-    try:
-        payload = jwt.decode(token.password, SECRET_KEY, algorithms=[ALGORITHM])
-        username: str = payload.get("sub")
-        if username is None:
-            raise credentials_exception
-        return {"username": username}
-    except JWTError:
-        raise credentials_exception
-
-# Routes
-@app.post("/api/submit")
-async def submit_query(name: str = Form(...), email: str = Form(...), message: str = Form(...)):
-    query = {"name": name, "email": email, "message": message, "created_at": datetime.utcnow()}
-    result = collection.insert_one(query)
-    return JSONResponse(content={"id": str(result.inserted_id), "notify": True})
-
-@app.post("/api/login")
-async def login(form_data: OAuth2PasswordRequestForm = Depends()):
-    user = authenticate_user(form_data.username, form_data.password)
-    if not user:
-        raise HTTPException(status_code=401, detail="Invalid credentials")
-    token = create_access_token({"sub": user["username"]})
-    return {"access_token": token, "token_type": "bearer"}
-
-@app.get("/api/queries")
-async def get_queries(token: str):
-    try:
-        jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-        queries = list(collection.find())
-        for q in queries:
-            q["_id"] = str(q["_id"])
-        return queries
-    except JWTError:
-        raise HTTPException(status_code=401, detail="Invalid token")
-
-@app.delete("/api/queries/{query_id}")
-async def delete_query(query_id: str, token: str):
-    try:
-        jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-        result = collection.delete_one({"_id": ObjectId(query_id)})
-        return {"deleted": result.deleted_count == 1}
-    except JWTError:
-        raise HTTPException(status_code=401, detail="Invalid token")
+from fastapi import FastAPI, HTTPException, Depends, Request, status, Form, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.security import OAuth2PasswordRequestForm
+from jose import JWTError, jwt
+from passlib.context import CryptContext
+from datetime import datetime, timedelta
+from pymongo import MongoClient
+from bson import ObjectId
+from fastapi.responses import JSONResponse
+from dotenv import load_dotenv
+import os
+from starlette.websockets import WebSocketState
+from jose import JWTError
+
+load_dotenv()
+
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: list[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)
+
+    async def broadcast(self, message: dict):
+        for connection in self.active_connections:
+            try:
+                await connection.send_json(message)
+            except Exception:
+                self.active_connections.remove(connection)
+
+manager = ConnectionManager()
+
+# Constants
+SECRET_KEY = os.environ.get("SECRET_KEY")
+ALGORITHM = os.environ.get("ALGORITHM")
+ACCESS_TOKEN_EXPIRE_MINUTES = 60
+
+# Admin credentials
+fake_admin_db = {
+    "admin": {
+        "username": os.environ.get("ADMIN_USERNAME"),
+        "hashed_password": os.environ.get("ADMIN_PASSWORD")
+    }
+}
+
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+client = MongoClient(os.environ.get("MONGODB_URI"))
+db = client["thehexatechdb"]
+collection = db["quotationsdb"]
+
+app = FastAPI()
+
+# CORS
+origins = ["*"]
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Auth Utilities
+def verify_password(plain, hashed):
+    return pwd_context.verify(plain, hashed)
+
+def get_password_hash(password):
+    return pwd_context.hash(password)
+
+def authenticate_user(username: str, password: str):
+    user = fake_admin_db.get(username)
+    if not user or not verify_password(password, user["hashed_password"]):
+        return False
+    return {"username": username}
+
+def create_access_token(data: dict, expires_delta=None):
+    to_encode = data.copy()
+    expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    to_encode.update({"exp": expire})
+    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+
+def get_current_user(token: str = Depends(OAuth2PasswordRequestForm)):
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Invalid credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    try:
+        payload = jwt.decode(token.password, SECRET_KEY, algorithms=[ALGORITHM])
+        username: str = payload.get("sub")
+        if username is None:
+            raise credentials_exception
+        return {"username": username}
+    except JWTError:
+        raise credentials_exception
+
+# Routes
+@app.post("/api/submit")
+async def submit_query(name: str = Form(...), email: str = Form(...), message: str = Form(...)):
+    query = {"name": name, "email": email, "message": message, "created_at": datetime.utcnow()}
+    result = collection.insert_one(query)
+    query["_id"] = str(result.inserted_id)
+
+    total_count = collection.count_documents({})
+    await manager.broadcast({
+        "event": "new_quote",
+        "data": {
+            "name": name,
+            "email": email,
+            "message": message,
+            "total_count": total_count
+        }
+    })
+    return JSONResponse(content={"id": query["_id"], "notify": True})
+
+
+@app.post("/api/login")
+async def login(form_data: OAuth2PasswordRequestForm = Depends()):
+    user = authenticate_user(form_data.username, form_data.password)
+    if not user:
+        raise HTTPException(status_code=401, detail="Invalid credentials")
+    token = create_access_token({"sub": user["username"], "role": "admin"})
+    return {"access_token": token, "token_type": "bearer"}
+
+@app.get("/api/queries")
+async def get_queries(token: str):
+    try:
+        jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        queries = list(collection.find())
+        for q in queries:
+            q["_id"] = str(q["_id"])
+        return queries
+    except JWTError:
+        raise HTTPException(status_code=401, detail="Invalid token")
+
+@app.delete("/api/queries/{query_id}")
+async def delete_query(query_id: str, token: str):
+    try:
+        jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        result = collection.delete_one({"_id": ObjectId(query_id)})
+        return {"deleted": result.deleted_count == 1}
+    except JWTError:
+        raise HTTPException(status_code=401, detail="Invalid token")
+
+
+@app.websocket("/ws/notifications")
+async def websocket_endpoint(websocket: WebSocket):
+    token = websocket.query_params.get("token")
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        username = payload.get("sub")
+        role = payload.get("role")
+        if not username or role != "admin":
+            await websocket.close(code=1008)
+            return
+    except JWTError:
+        await websocket.close(code=1008)
+        return
+
+    await manager.connect(websocket)
+    print(f"[WS CONNECT] {username} (admin) connected at {datetime.utcnow()} from {websocket.client.host}")
+
+    try:
+        while True:
+            await websocket.receive_text()
+    except WebSocketDisconnect:
+        print(f"[WS DISCONNECT] {username} disconnected at {datetime.utcnow()}")
+        manager.disconnect(websocket)