31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
from functools import wraps
|
|
from flask import session, flash, redirect, url_for
|
|
from db import query_db
|
|
|
|
def login_required(f):
|
|
"""Require login for route"""
|
|
@wraps(f)
|
|
def decorated_function(*args, **kwargs):
|
|
if 'user_id' not in session:
|
|
flash('Please log in to access this page', 'warning')
|
|
return redirect(url_for('login'))
|
|
return f(*args, **kwargs)
|
|
return decorated_function
|
|
|
|
def role_required(*roles):
|
|
"""Require specific role(s) for route"""
|
|
def decorator(f):
|
|
@wraps(f)
|
|
def decorated_function(*args, **kwargs):
|
|
if 'user_id' not in session:
|
|
flash('Please log in to access this page', 'warning')
|
|
return redirect(url_for('login'))
|
|
|
|
user = query_db('SELECT role FROM Users WHERE user_id = ?', [session['user_id']], one=True)
|
|
if not user or user['role'] not in roles:
|
|
flash('You do not have permission to access this page', 'danger')
|
|
return redirect(url_for('dashboard'))
|
|
|
|
return f(*args, **kwargs)
|
|
return decorated_function
|
|
return decorator |