Add decorators db_lookup and protect_csrf

csrf.py

+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

database.py

+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