From 8561391cd2ac2dab582a55313bf2d682b9184473 Mon Sep 17 00:00:00 2001
From: Adam Shimali <adam.shimali@gmail.com>
Date: Mon, 14 Mar 2016 16:30:48 +0000
Subject: [PATCH] The verify view was not passing along the next param to the
 two factor view.

Now if it is passed and it is a url on the same domain that request
originates from then it is used.
---
 app/main/views/sign_in.py               |  8 +++--
 app/main/views/two_factor.py            | 17 +++++++++-
 tests/app/main/views/test_two_factor.py | 44 +++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/app/main/views/sign_in.py b/app/main/views/sign_in.py
index 5b598e4ae..9121b3b9b 100644
--- a/app/main/views/sign_in.py
+++ b/app/main/views/sign_in.py
@@ -4,7 +4,8 @@ from flask import (
     url_for,
     session,
     abort,
-    flash
+    flash,
+    request
 )
 
 from flask.ext.login import (current_user, login_fresh, confirm_login)
@@ -41,7 +42,10 @@ def sign_in():
                 return redirect(url_for('.verify'))
             elif user.is_active():
                 users_dao.send_verify_code(user.id, 'sms', user.mobile_number)
-                return redirect(url_for('.two_factor'))
+                if request.args.get('next'):
+                    return redirect(url_for('.two_factor', next=request.args.get('next')))
+                else:
+                    return redirect(url_for('.two_factor'))
         # Vague error message for login in case of user not known, locked, inactive or password not verified
         flash('Username or password is incorrect')
 
diff --git a/app/main/views/two_factor.py b/app/main/views/two_factor.py
index 37618ed0d..bc6a2e945 100644
--- a/app/main/views/two_factor.py
+++ b/app/main/views/two_factor.py
@@ -3,7 +3,8 @@ from flask import (
     render_template,
     redirect,
     session,
-    url_for
+    url_for,
+    request
 )
 
 from flask_login import login_user
@@ -37,9 +38,23 @@ def two_factor():
             login_user(user, remember=True)
         finally:
             del session['user_details']
+
+        next_url = request.args.get('next')
+        if next_url and _is_safe_redirect_url(next_url):
+            return redirect(next_url)
+
         if len(services) == 1:
             return redirect(url_for('main.service_dashboard', service_id=services[0]['id']))
         else:
             return redirect(url_for('main.choose_service'))
 
     return render_template('views/two-factor.html', form=form)
+
+
+# see http://flask.pocoo.org/snippets/62/
+def _is_safe_redirect_url(target):
+    from urllib.parse import urlparse, urljoin
+    host_url = urlparse(request.host_url)
+    redirect_url = urlparse(urljoin(request.host_url, target))
+    return redirect_url.scheme in ('http', 'https') and \
+        host_url.netloc == redirect_url.netloc
diff --git a/tests/app/main/views/test_two_factor.py b/tests/app/main/views/test_two_factor.py
index 51d49324d..fb7cf7fb7 100644
--- a/tests/app/main/views/test_two_factor.py
+++ b/tests/app/main/views/test_two_factor.py
@@ -40,6 +40,50 @@ def test_should_login_user_and_redirect_to_service_dashboard(app_,
             )
 
 
+def test_should_login_user_and_should_redirect_to_next_url(app_,
+                                                           api_user_active,
+                                                           mock_get_user,
+                                                           mock_get_user_by_email,
+                                                           mock_check_verify_code,
+                                                           mock_get_services):
+    with app_.test_request_context():
+        with app_.test_client() as client:
+            with client.session_transaction() as session:
+                session['user_details'] = {
+                    'id': api_user_active.id,
+                    'email': api_user_active.email_address}
+            response = client.post(url_for('main.two_factor', next='/services/{}/dashboard'.format(SERVICE_ONE_ID)),
+                                   data={'sms_code': '12345'})
+            assert response.status_code == 302
+            assert response.location == url_for(
+                'main.service_dashboard',
+                service_id=SERVICE_ONE_ID,
+                _external=True
+            )
+
+
+def test_should_login_user_and_not_redirect_to_external_url(app_,
+                                                            api_user_active,
+                                                            mock_get_user,
+                                                            mock_get_user_by_email,
+                                                            mock_check_verify_code,
+                                                            mock_get_services_with_one_service):
+    with app_.test_request_context():
+        with app_.test_client() as client:
+            with client.session_transaction() as session:
+                session['user_details'] = {
+                    'id': api_user_active.id,
+                    'email': api_user_active.email_address}
+            response = client.post(url_for('main.two_factor', next='http://www.google.com'),
+                                   data={'sms_code': '12345'})
+            assert response.status_code == 302
+            assert response.location == url_for(
+                'main.service_dashboard',
+                service_id=SERVICE_ONE_ID,
+                _external=True
+            )
+
+
 def test_should_login_user_and_redirect_to_choose_services(app_,
                                                            api_user_active,
                                                            mock_get_user,