Added return_url to admin login

3 files changed, 10 insertions, 5 deletions

diff --git a/app.py b/app.py
index ae6ec36..6b3e4da 100644
--- a/app.py
+++ b/app.py
@@ -38,13 +38,14 @@ def inject():
     return {
         'projects': from_json(db.all(), List[Project]),
         'empty_project': Project(),
-        'admin': session.get('logged_in', False)
+        'admin': session.get('logged_in', False),
+        'no_login': False,
     }
 
 
 @app.route('/admin/login/', methods=['GET'])
 def login():
-    return render_template('login.html')
+    return render_template('login.html', return_url=request.args.get('return_url', '/'), no_login=True)
 
 
 @app.route('/admin/login', methods=['POST'])
@@ -52,7 +53,7 @@ def check_login():
     if not check_password(request.form.get('pass')):
         return redirect(url_for('login'))
     session['logged_in'] = True
-    return redirect('/')
+    return redirect(request.args.get('return_url', '/'))
 
 
 @app.route('/admin/logout')
diff --git a/templates/base.html b/templates/base.html
index 9f45c87..c48d1c3 100644
--- a/templates/base.html
+++ b/templates/base.html
@@ -25,7 +25,11 @@
         <a href="{{ url_for('new_project_form') }}">New Project</a>
         {% block admin_footer %}{% endblock %}
     {% else %}
-        <a href="{{ url_for('login') }}">Admin login</a>
+        {% if no_login %}
+            <a href="{{ return_url }}">Go back</a>
+        {% else %}
+            <a href="{{ url_for('login') }}?return_url={{ request.path }}">Admin login</a>
+        {% endif %}
     {% endif %}
     {% block footer %}
     {% endblock %}
diff --git a/templates/login.html b/templates/login.html
index 32059c6..14a22e9 100644
--- a/templates/login.html
+++ b/templates/login.html
@@ -1,6 +1,6 @@
 {% extends "base.html" %}
 {% block body %}
-    <form method="post" action="{{ url_for('check_login') }}">
+    <form method="post" action="{{ url_for('check_login') }}?return_url={{ return_url }}">
         <p style="color: red;">Admin access only</p>
         <p>
             <label for="user">Username</label>