Commit c7eac74c authored by Robin Sonnabend's avatar Robin Sonnabend

Add decorators db_lookup and protect_csrf

parent 3dfc2b71
from functools import wraps
from hmac import compare_digest
from flask import request, abort, session
def get_csrf_token():
if "_csrf" not in session:
session["_csrf"] = hashlib.sha1(os.urandom(64)).hexdigest()
return session["_csrf"]
def protect_csrf(function):
@wraps(function)
def _decorated_function(*args, **kwargs):
token = request.args.get("csrf_token")
true_token = get_csrf_token()
if token is None or not compare_digest(token, true_token):
abort(400)
return function(*args, **kwargs)
return _decorated_function
from flask import flash
from functools import wraps
from . import back
ID_KEY = "id"
KEY_NOT_PRESENT_MESSAGE = "Missing {}_id."
OBJECT_DOES_NOT_EXIST_MESSAGE = "There is no {} with id {}."
def default_redirect():
return back.redirect()
def login_redirect():
return back.redirect("login")
def db_lookup(*models, check_exists=True):
def _decorator(function):
@wraps(function)
def _decorated_function(*args, **kwargs):
for model in models:
key = model.__model_name__
id_key = "{}_{}".format(key, ID_KEY)
if id_key not in kwargs:
flash(KEY_NOT_PRESENT_MESSAGE.format(key), "alert-error")
return default_redirect()
obj_id = kwargs[id_key]
obj = model.query.filter_by(id=obj_id).first()
if check_exists and obj is None:
model_name = model.__class__.__name__
flash(OBJECT_DOES_NOT_EXIST_MESSAGE.format(
model_name, obj_id),
"alert-error")
return default_redirect()
kwargs[key] = obj
kwargs.pop(id_key)
return function(*args, **kwargs)
return _decorated_function
return _decorator
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment