diff --git a/csrf.py b/csrf.py new file mode 100644 index 0000000000000000000000000000000000000000..a415163420db2baa8121ee5a10323dc6008a62e7 --- /dev/null +++ b/csrf.py @@ -0,0 +1,21 @@ +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 diff --git a/database.py b/database.py new file mode 100644 index 0000000000000000000000000000000000000000..f6cbc567ec2d1b3abe6b8286cd5ebbd17ee600b3 --- /dev/null +++ b/database.py @@ -0,0 +1,42 @@ +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