82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
"""
|
||
认证装饰器与 session 校验
|
||
"""
|
||
from functools import wraps
|
||
from flask import request, redirect, url_for, session, jsonify
|
||
|
||
from utils.db import get_db
|
||
|
||
|
||
def is_session_user_valid():
|
||
"""校验 session 中的 user_id 是否在数据库中仍存在;不存在则清除 session 并返回 False"""
|
||
uid = session.get('user_id')
|
||
if not uid:
|
||
return False
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT id FROM users WHERE id = %s", (uid,))
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row:
|
||
session.clear()
|
||
return False
|
||
return True
|
||
except Exception:
|
||
session.clear()
|
||
return False
|
||
|
||
|
||
def get_current_admin_role():
|
||
"""获取当前登录用户的管理角色:super_admin / admin / None(非管理员)"""
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"SELECT id, username, is_admin, role, created_by_id FROM users WHERE id = %s",
|
||
(session['user_id'],)
|
||
)
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row or not row.get('is_admin'):
|
||
return None, None
|
||
return row.get('role') or ('super_admin' if row.get('created_by_id') is None else 'admin'), row
|
||
except Exception:
|
||
return None, None
|
||
|
||
|
||
def login_required(f):
|
||
@wraps(f)
|
||
def decorated(*args, **kwargs):
|
||
if not session.get('user_id') or not is_session_user_valid():
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
|
||
return jsonify({'success': False, 'error': '未登录'}), 401
|
||
return redirect(url_for('auth.login'))
|
||
return f(*args, **kwargs)
|
||
return decorated
|
||
|
||
|
||
def admin_required(f):
|
||
@wraps(f)
|
||
def decorated(*args, **kwargs):
|
||
if not session.get('user_id'):
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
|
||
return jsonify({'success': False, 'error': '未登录'}), 401
|
||
return redirect(url_for('auth.login'))
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT is_admin, role FROM users WHERE id = %s", (session['user_id'],))
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row or not row.get('is_admin'):
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
|
||
return jsonify({'success': False, 'error': '需要管理员权限'}), 403
|
||
return redirect(url_for('main.admin_page'))
|
||
except Exception as e:
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
return redirect(url_for('main.admin_page'))
|
||
return f(*args, **kwargs)
|
||
return decorated
|