Files

189 lines
6.8 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from typing import Any
from asgiref.sync import sync_to_async
from starlette.background import BackgroundTask
from starlette.datastructures import UploadFile
from starlette.requests import Request
from starlette.types import ASGIApp, Scope, Receive, Send
from backend.app.common.enums import OperaLogCipherType
from backend.app.common.log import log
from backend.app.core.conf import settings
from backend.app.schemas.opera_log import CreateOperaLog
from backend.app.services.opera_log_service import OperaLogService
from backend.app.utils.encrypt import AESCipher, Md5Cipher
from backend.app.utils.request_parse import parse_user_agent_info, parse_ip_info
from backend.app.utils.timezone import timezone_utils
class OperaLogMiddleware:
"""操作日志中间件"""
def __init__(self, app: ASGIApp):
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope['type'] != 'http':
await self.app(scope, receive, send)
return
request = Request(scope=scope, receive=receive)
# 排除记录白名单
path = request.url.path
if path in settings.OPERA_LOG_EXCLUDE or not path.startswith(f'{settings.API_V1_STR}'):
await self.app(scope, receive, send)
return
# 请求信息解析
user_agent, device, os, browser = await parse_user_agent_info(request)
ip, country, region, city = await parse_ip_info(request)
try:
# 此信息依赖于 jwt 中间件
username = request.user.username
except AttributeError:
username = None
method = request.method
args = await self.get_request_args(request)
# 设置附加请求信息
request.state.ip = ip
request.state.country = country
request.state.region = region
request.state.city = city
request.state.user_agent = user_agent
request.state.os = os
request.state.browser = browser
request.state.device = device
# 执行请求
start_time = timezone_utils.get_timezone_datetime()
code, msg, status, err = await self.execute_request(request, send)
end_time = timezone_utils.get_timezone_datetime()
cost_time = (end_time - start_time).total_seconds() * 1000.0
router = request.scope.get('route')
summary = getattr(router, 'summary') or ''
args.update(request.path_params)
# 脱敏处理
args = await self.desensitization(args)
# 日志创建
opera_log_in = CreateOperaLog(
username=username,
method=method,
title=summary,
path=path,
ip=ip,
country=country,
region=region,
city=city,
user_agent=user_agent,
os=os,
browser=browser,
device=device,
args=args,
status=status,
code=code,
msg=msg,
cost_time=cost_time,
opera_time=start_time,
)
back = BackgroundTask(OperaLogService.create, obj_in=opera_log_in)
await back()
# 错误抛出
if err:
raise err from None
async def execute_request(self, request: Request, send: Send) -> tuple:
err: Any = None
try:
# 详见 https://github.com/tiangolo/fastapi/discussions/8385#discussioncomment-6117967
async def wrapped_rcv_gen():
async for _ in request.stream():
yield {'type': 'http.request', 'body': await request.body()}
async for message in request.receive: # type: ignore
# assert message['type'] == 'http.response.body'
# body = message.get('body', b'')
# if body:
# yield body
yield message
wrapped_rcv = wrapped_rcv_gen().__anext__
await self.app(request.scope, wrapped_rcv, send)
code, msg, status = await self.exception_middleware_handler(request)
except Exception as e:
log.exception(e)
code = getattr(e, 'code', None) or '500'
msg = getattr(e, 'msg', None) or 'Internal Server Error'
status = 0
err = e
return code, msg, status, err
@staticmethod
@sync_to_async
def exception_middleware_handler(request: Request):
# 预置响应信息
code: str = '200'
msg: str = 'Success'
status: int = 1
try:
http_exception = request.state.__request_http_exception__
except AttributeError:
pass
else:
code = http_exception.get('code', '500')
msg = http_exception.get('msg', 'Internal Server Error')
status = 0
try:
validation_exception = request.state.__request_validation_exception__
except AttributeError:
pass
else:
code = validation_exception.get('code', '400')
msg = validation_exception.get('msg', 'Bad Request')
status = 0
return code, msg, status
@staticmethod
async def get_request_args(request: Request) -> dict:
args = dict(request.query_params)
form_data = await request.form()
if len(form_data) > 0:
args.update({k: v.filename if isinstance(v, UploadFile) else v for k, v in form_data.items()})
else:
body_data = await request.body()
if body_data:
json_data = await request.json()
args.update(json_data)
return args
@staticmethod
@sync_to_async
def desensitization(args: dict):
if len(args) > 0:
match settings.OPERA_LOG_ENCRYPT:
case OperaLogCipherType.aes:
for key in args.keys():
if key in settings.OPERA_LOG_ENCRYPT_INCLUDE:
args[key] = (
AESCipher(settings.OPERA_LOG_ENCRYPT_SECRET_KEY).encrypt(
bytes(args[key], encoding='utf-8')
)
).hex()
case OperaLogCipherType.md5:
for key in args.keys():
if key in settings.OPERA_LOG_ENCRYPT_INCLUDE:
args[key] = Md5Cipher.encrypt(args[key])
case OperaLogCipherType.plan:
pass
case _:
for key in args.keys():
if key in settings.OPERA_LOG_ENCRYPT_INCLUDE:
args[key] = '******'
return args if len(args) > 0 else None