38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
from functools import wraps
|
||
from flask import session, jsonify, redirect, url_for, request
|
||
from models import Admin
|
||
import pyotp
|
||
|
||
def admin_required(f):
|
||
"""管理员权限装饰器"""
|
||
@wraps(f)
|
||
def decorated_function(*args, **kwargs):
|
||
if 'admin_id' not in session:
|
||
# 如果是API请求，返回JSON
|
||
if request.path.startswith('/admin/api/'):
|
||
return jsonify({'success': False, 'message': '请先登录'}), 401
|
||
# 如果是页面请求，重定向到登录页
|
||
return redirect(url_for('admin.login'))
|
||
return f(*args, **kwargs)
|
||
return decorated_function
|
||
|
||
def verify_2fa(admin: Admin, code: str) -> bool:
|
||
"""验证2FA代码"""
|
||
if not admin.is_2fa_enabled or not admin.totp_secret:
|
||
return True
|
||
|
||
totp = pyotp.TOTP(admin.totp_secret)
|
||
return totp.verify(code, valid_window=1)
|
||
|
||
def generate_2fa_secret() -> str:
|
||
"""生成2FA密钥"""
|
||
return pyotp.random_base32()
|
||
|
||
def get_2fa_qrcode_url(admin: Admin, secret: str) -> str:
|
||
"""获取2FA二维码URL"""
|
||
totp = pyotp.TOTP(secret)
|
||
return totp.provisioning_uri(
|
||
name=admin.username,
|
||
issuer_name='短视频解析平台'
|
||
)
|