Authentication in Flask Apps: Making a 'login required' decorator (+13, -2)

__init__.py (+13, -2)

From: curriculum/section13/lectures/05_login_required_decorator/start/app/__init__.py

To: curriculum/section13/lectures/05_login_required_decorator/end/app/__init__.py

            
            index 9fa838e..b4ce994 100644
--- a/curriculum/section13/lectures/05_login_required_decorator/start/app/__init__.py
+++ b/curriculum/section13/lectures/05_login_required_decorator/end/app/__init__.py
@@ -1,4 +1,5 @@
 import os
+import functools
 from flask import (
     Flask,
     session,
@@ -18,15 +19,25 @@ app.secret_key = "lkaQT-kAb6aIvqWETVcCQ28F-j-rP_PSEaCDdTynkXA"
 users = {}
 
 
+def login_required(route):
+    @functools.wraps(route)
+    def route_wrapper(*args, **kwargs):
+        if session.get("email") is None:
+            return redirect(url_for("login"))
+
+        return route(*args, **kwargs)
+
+    return route_wrapper
+
+
 @app.get("/")
 def home():
     return render_template("home.html", email=session.get("email"))
 
 
 @app.get("/protected")
+@login_required
 def protected():
-    if not session.get("email"):
-        abort(401)
     return render_template("protected.html")