auth pages less irritating

2 files changed, 25 insertions, 6 deletions

diff --git a/apioforum/auth.py b/apioforum/auth.py
index 49808ea..547f15e 100644
--- a/apioforum/auth.py
+++ b/apioforum/auth.py
@@ -8,6 +8,9 @@ import functools
 
 bp = Blueprint("auth", __name__, url_prefix="/auth")
 
+def get_next():
+    return request.args.get('next',url_for('index'))
+
 @bp.route("/login",methods=('GET','POST'))
 def login():
     if request.method == "POST":
@@ -28,12 +31,14 @@ def login():
         if err is None:
             session.clear()
             session['user'] = username
-            return redirect(url_for('auth.cool'))
+            flash("logged in successfully")
+            return redirect(get_next())
 
         flash(err)
         
     return render_template("auth/login.html")
 
+
 @bp.route("/register", methods=("GET","POST"))
 def register():
     if request.method == "POST":
@@ -58,7 +63,8 @@ def register():
             db.commit()
             flash("successfully created account")
             session['user'] = username
-            return redirect(url_for("auth.cool"))
+            flash("registered successfully")
+            return redirect(get_next())
 
         flash(err)
             
@@ -67,7 +73,8 @@ def register():
 @bp.route("/logout")
 def logout():
     session.clear()
-    return redirect(url_for("auth.cool"))
+    flash("logged out successfully")
+    return redirect(get_next())
 
 @bp.before_app_request
 def load_user():
diff --git a/apioforum/templates/base.html b/apioforum/templates/base.html
index 847b8b2..f81f413 100644
--- a/apioforum/templates/base.html
+++ b/apioforum/templates/base.html
@@ -17,10 +17,22 @@
 
                     {% if g.user %}
                     <li>{{ g.user }}</li>
-                    <li><a href="{{ url_for('auth.logout') }}">logout</a></li>
+                    <li>
+                        <a href="{{ url_for('auth.logout',next=request.path) }}">
+                            logout
+                        </a>
+                    </li>
                     {% else %}
-                    <li><a href="{{ url_for('auth.login') }}">login</a></li>
-                    <li><a href="{{ url_for('auth.register') }}">register</a></li>
+                    <li>
+                        <a href="{{ url_for('auth.login',next=request.path) }}">
+                            login
+                        </a>
+                    </li>
+                    <li>
+                        <a href="{{ url_for('auth.register',next=request.path) }}">
+                            register
+                        </a>
+                    </li>
                     {% endif %}
                 </ul>
             </nav>