109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
import json
|
|
import logging
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse, Response
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from starlette.middleware.cors import CORSMiddleware
|
|
|
|
from app.core.config import get_settings
|
|
from app.core.crypto import decrypt_json, encrypt_json
|
|
|
|
logger = logging.getLogger("gsfdc.access")
|
|
|
|
_SKIP_ENCRYPT_PREFIXES = ("/docs", "/redoc", "/openapi.json")
|
|
|
|
|
|
def _skip_encrypt(path: str, content_type: str) -> bool:
|
|
if path.startswith(_SKIP_ENCRYPT_PREFIXES):
|
|
return True
|
|
if "multipart/form-data" in content_type:
|
|
return True
|
|
return False
|
|
|
|
|
|
class AccessGuardMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
client = request.client.host if request.client else ""
|
|
if _ip_blocked(client):
|
|
return JSONResponse(status_code=403, content={"detail": "访问被拒绝"})
|
|
|
|
if get_settings().api_payload_encrypt and not _skip_encrypt(
|
|
request.url.path, request.headers.get("content-type", "")
|
|
):
|
|
request = await _decrypt_request(request)
|
|
|
|
response = await call_next(request)
|
|
logger.info("%s %s %s", request.method, request.url.path, response.status_code)
|
|
_record_admin_call(request, response.status_code)
|
|
|
|
if get_settings().api_payload_encrypt and not _skip_encrypt(
|
|
request.url.path, response.headers.get("content-type", "")
|
|
):
|
|
return await _encrypt_response(response)
|
|
return response
|
|
|
|
|
|
async def _decrypt_request(request: Request) -> Request:
|
|
body = await request.body()
|
|
raw = body
|
|
if body:
|
|
try:
|
|
envelope = json.loads(body)
|
|
except json.JSONDecodeError:
|
|
envelope = None
|
|
if isinstance(envelope, dict) and "data" in envelope:
|
|
plain = decrypt_json(envelope)
|
|
raw = json.dumps(plain, ensure_ascii=False).encode()
|
|
|
|
async def receive():
|
|
return {"type": "http.request", "body": raw, "more_body": False}
|
|
|
|
return Request(request.scope, receive)
|
|
|
|
|
|
async def _encrypt_response(response: Response) -> Response:
|
|
body = b""
|
|
async for chunk in response.body_iterator:
|
|
body += chunk
|
|
if not body:
|
|
return response
|
|
try:
|
|
payload = json.loads(body)
|
|
except json.JSONDecodeError:
|
|
return Response(content=body, status_code=response.status_code, headers=dict(response.headers))
|
|
encrypted = encrypt_json(payload if isinstance(payload, (dict, list)) else {"data": payload})
|
|
return JSONResponse(status_code=response.status_code, content=encrypted)
|
|
|
|
|
|
def _ip_blocked(ip: str) -> bool:
|
|
if not ip:
|
|
return False
|
|
try:
|
|
from app.services.ip_guard import is_ip_blocked
|
|
|
|
return is_ip_blocked(ip)
|
|
except Exception:
|
|
return ip in get_settings().blocked_ips
|
|
|
|
|
|
def _record_admin_call(request: Request, status_code: int) -> None:
|
|
try:
|
|
from app.services.audit import record_admin_call
|
|
|
|
record_admin_call(request, status_code)
|
|
except Exception:
|
|
logger.warning("admin audit hook failed")
|
|
|
|
|
|
def install_middleware(app: FastAPI) -> None:
|
|
settings = get_settings()
|
|
app.add_middleware(AccessGuardMiddleware)
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=settings.cors_origin_list,
|
|
allow_credentials=False,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|