Secure back redirects with WTForms.

3 files changed, 50 insertions, 11 deletions

diff --git a/pyaggr3g470r/forms.py b/pyaggr3g470r/forms.py
index 77799c4d..0998c2e6 100644
--- a/pyaggr3g470r/forms.py
+++ b/pyaggr3g470r/forms.py
@@ -26,14 +26,16 @@ __revision__ = "$Date: 2015/05/06 $"
 __copyright__ = "Copyright (c) Cedric Bonhomme"
 __license__ = "GPLv3"
 
-from flask import flash
+
+from flask import flash, request, url_for, redirect
 from flask.ext.wtf import Form
 from flask.ext.babel import lazy_gettext
 from wtforms import TextField, TextAreaField, PasswordField, BooleanField, \
-        SubmitField, IntegerField, validators
+        SubmitField, IntegerField, validators, HiddenField
 from flask.ext.wtf.html5 import EmailField
 from flask_wtf import RecaptchaField
 
+from pyaggr3g470r import utils
 from pyaggr3g470r.models import User
 
 class SignupForm(Form):
@@ -59,8 +61,24 @@ class SignupForm(Form):
             validated = False
         return validated
 
+class RedirectForm(Form):
+    """
+    Secure back redirects with WTForms.
+    """
+    next = HiddenField()
+
+    def __init__(self, *args, **kwargs):
+        Form.__init__(self, *args, **kwargs)
+        if not self.next.data:
+            self.next.data = utils.get_redirect_target() or ''
+
+    def redirect(self, endpoint='home', **values):
+        if utils.is_safe_url(self.next.data):
+            return redirect(self.next.data)
+        target = utils.get_redirect_target()
+        return redirect(target or url_for(endpoint, **values))
 
-class SigninForm(Form):
+class SigninForm(RedirectForm):
     """
     Sign in form (connection to pyAggr3g470r).
     """
diff --git a/pyaggr3g470r/utils.py b/pyaggr3g470r/utils.py
index 3d8bb483..bcea5109 100755
--- a/pyaggr3g470r/utils.py
+++ b/pyaggr3g470r/utils.py
@@ -49,11 +49,12 @@ import sqlalchemy
 try:
     from urlparse import urlparse, parse_qs, urlunparse
 except:
-    from urllib.parse import urlparse, parse_qs, urlunparse
+    from urllib.parse import urlparse, parse_qs, urlunparse, urljoin
 from bs4 import BeautifulSoup
 from datetime import timedelta
 from collections import Counter
 from contextlib import contextmanager
+from flask import request
 
 import conf
 from flask import g
@@ -65,6 +66,25 @@ logger = logging.getLogger(__name__)
 
 ALLOWED_EXTENSIONS = set(['xml', 'opml', 'json'])
 
+def is_safe_url(target):
+    """
+    Ensures that a redirect target will lead to the same server.
+    """
+    ref_url = urlparse(request.host_url)
+    test_url = urlparse(urljoin(request.host_url, target))
+    return test_url.scheme in ('http', 'https') and \
+           ref_url.netloc == test_url.netloc
+
+def get_redirect_target():
+    """
+    Looks at various hints to find the redirect target.
+    """
+    for target in request.args.get('next'), request.referrer:
+        if not target:
+            continue
+        if is_safe_url(target):
+            return target
+
 def allowed_file(filename):
     """
     Check if the uploaded file is allowed.
diff --git a/pyaggr3g470r/views/views.py b/pyaggr3g470r/views/views.py
index 29b865e0..69c2b50b 100644
--- a/pyaggr3g470r/views/views.py
+++ b/pyaggr3g470r/views/views.py
@@ -38,7 +38,8 @@ from bootstrap import application as app, db
 from flask import render_template, request, flash, session, \
                   url_for, redirect, g, current_app, make_response
 from flask.ext.login import LoginManager, login_user, logout_user, \
-                            login_required, current_user, AnonymousUserMixin
+                            login_required, current_user, AnonymousUserMixin, \
+                            login_url
 from flask.ext.principal import Principal, Identity, AnonymousIdentity, \
                                 identity_changed, identity_loaded, Permission,\
                                 RoleNeed, UserNeed
@@ -65,6 +66,10 @@ admin_permission = Permission(RoleNeed('admin'))
 
 login_manager = LoginManager()
 login_manager.init_app(app)
+login_manager.login_message = gettext('Authentication required.')
+login_manager.login_message_category = "info"
+login_manager.login_view = 'login'
+
 logger = logging.getLogger(__name__)
 
 #
@@ -98,7 +103,6 @@ def load_user(id):
     # Return an instance of the User model
     return UserController().get(id=id)
 
-
 #
 # Custom error pages.
 #
@@ -110,7 +114,7 @@ def authentication_required(e):
 @app.errorhandler(403)
 def authentication_failed(e):
     flash(gettext('Forbidden.'), 'danger')
-    return redirect(url_for('home'))
+    return redirect(url_for('login'))
 
 @app.errorhandler(404)
 def page_not_found(e):
@@ -151,10 +155,8 @@ def login():
     """
     if g.user is not None and g.user.is_authenticated():
         return redirect(url_for('home'))
-
     g.user = AnonymousUserMixin()
     form = SigninForm()
-
     if form.validate_on_submit():
         user = UserController().get(email=form.email.data)
         login_user(user)
@@ -162,10 +164,9 @@ def login():
         session['email'] = form.email.data
         identity_changed.send(current_app._get_current_object(),
                               identity=Identity(user.id))
-        return redirect(url_for('home'))
+        return form.redirect('home')
     return render_template('login.html', form=form)
 
-
 @app.route('/logout')
 @login_required
 def logout():