76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
from collections.abc import Awaitable, Callable
|
|
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
from sqlalchemy import text
|
|
|
|
try:
|
|
# For module mode: `uvicorn app.main:app`
|
|
from app.db import SessionLocal
|
|
from app.security import auth_error_response, decode_access_token
|
|
except ModuleNotFoundError:
|
|
# For script mode: `python app/main.py` or VS Code "Run Python File"
|
|
from db import SessionLocal
|
|
from security import auth_error_response, decode_access_token
|
|
|
|
|
|
EXCLUDED_PATH_PREFIXES = (
|
|
"/health",
|
|
"/docs",
|
|
"/redoc",
|
|
"/openapi.json",
|
|
"/auth/login",
|
|
)
|
|
|
|
|
|
def _is_excluded_path(path: str) -> bool:
|
|
for prefix in EXCLUDED_PATH_PREFIXES:
|
|
if path == prefix or path.startswith(f"{prefix}/"):
|
|
return True
|
|
return False
|
|
|
|
|
|
def install_auth_middleware(app: FastAPI) -> None:
|
|
@app.middleware("http")
|
|
async def auth_middleware(
|
|
request: Request,
|
|
call_next: Callable[[Request], Awaitable],
|
|
):
|
|
path = request.url.path
|
|
if request.method == "OPTIONS" or _is_excluded_path(path):
|
|
return await call_next(request)
|
|
|
|
auth_header = request.headers.get("Authorization")
|
|
if not auth_header:
|
|
return auth_error_response("missing authorization header")
|
|
|
|
parts = auth_header.split(" ", 1)
|
|
if len(parts) != 2 or parts[0].lower() != "bearer":
|
|
return auth_error_response("invalid authorization format")
|
|
|
|
try:
|
|
user_id = decode_access_token(parts[1].strip())
|
|
except HTTPException:
|
|
return auth_error_response("invalid or expired access token")
|
|
|
|
with SessionLocal() as db:
|
|
user_row = (
|
|
db.execute(
|
|
text(
|
|
"""
|
|
SELECT id, status
|
|
FROM chat_user
|
|
WHERE id = :user_id
|
|
LIMIT 1
|
|
"""
|
|
),
|
|
{"user_id": user_id},
|
|
)
|
|
.mappings()
|
|
.first()
|
|
)
|
|
if not user_row or int(user_row["status"]) != 1:
|
|
return auth_error_response("user not found or disabled")
|
|
|
|
request.state.user_id = user_id
|
|
return await call_next(request)
|