57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
from fastapi import Request
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.responses import JSONResponse
|
|
from sqlalchemy.orm import Session
|
|
from app.core.database.db_session import SessionLocal
|
|
from app.models.SettingModel import Setting
|
|
from jose import jwt, JWTError
|
|
from app.core.config.Config import settings, load_public_key_bytes
|
|
|
|
MASTER_ADMIN_EMAIL = "adithiyan.elan@gmail.com"
|
|
|
|
class KillSwitchMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
# Allow docs and OpenAPI schemas to always bypass
|
|
if request.url.path in ["/docs", "/redoc", "/openapi.json"]:
|
|
return await call_next(request)
|
|
|
|
# Open db session safely
|
|
db: Session = SessionLocal()
|
|
global_disable = False
|
|
try:
|
|
# Query the GLOBAL_DISABLE setting key
|
|
item = db.query(Setting).filter(Setting.setting_key == "GLOBAL_DISABLE").first()
|
|
if item:
|
|
# Expecting value like {"enabled": true}
|
|
val = item.setting_value
|
|
if isinstance(val, dict) and val.get("enabled") is True:
|
|
global_disable = True
|
|
except Exception:
|
|
# If database lookup fails, fail secure or run normally?
|
|
# Let's run normally to prevent site breakdown during DB connection blips,
|
|
# but log the failure in real setups.
|
|
pass
|
|
finally:
|
|
db.close()
|
|
|
|
if not global_disable:
|
|
return await call_next(request)
|
|
|
|
# Bypass check for MASTER ADMIN
|
|
auth_header = request.headers.get("Authorization")
|
|
if auth_header and auth_header.startswith("Bearer "):
|
|
token = auth_header.split(" ")[1]
|
|
try:
|
|
public_key = load_public_key_bytes()
|
|
payload = jwt.decode(token, public_key, algorithms=[settings.ALGORITHM])
|
|
email = payload.get("email")
|
|
if email == MASTER_ADMIN_EMAIL:
|
|
# Let master bypass the lock
|
|
return await call_next(request)
|
|
except JWTError:
|
|
pass
|
|
|
|
return JSONResponse(
|
|
status_code=403,
|
|
content={"detail": "SYSTEM DISABLED BY EMERGENCY KILL SWITCH"}
|
|
)
|