From c946f85f9d4474d041e83d9ef774b3ff4d4d94da Mon Sep 17 00:00:00 2001
From: Rebecca Law <rebecca.law@digital.cabinet-office.gov.uk>
Date: Tue, 8 Dec 2015 09:21:51 +0000
Subject: [PATCH] 109638656: Send sms code from sign-in post.

---
 app/main/views/__init__.py              | 30 +++++++++++++++++++++++++
 app/main/views/register.py              | 27 +---------------------
 app/main/views/sign_in.py               |  8 +++++--
 app/main/views/two_factor.py            |  2 ++
 tests/app/main/views/test_two_factor.py |  4 ++--
 5 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/app/main/views/__init__.py b/app/main/views/__init__.py
index e69de29bb..0b8514dda 100644
--- a/app/main/views/__init__.py
+++ b/app/main/views/__init__.py
@@ -0,0 +1,30 @@
+from random import randint
+from app import admin_api_client
+from app.main.exceptions import AdminApiClientException
+
+
+def create_verify_code():
+    return ''.join(["%s" % randint(0, 9) for _ in range(0, 5)])
+
+
+def send_sms_code(mobile_number):
+    sms_code = create_verify_code()
+    try:
+        admin_api_client.send_sms(mobile_number, message=sms_code, token=admin_api_client.auth_token)
+    except:
+        raise AdminApiClientException('Exception when sending sms.')
+    return sms_code
+
+
+def send_email_code(email):
+    email_code = create_verify_code()
+    try:
+        admin_api_client.send_email(email_address=email,
+                                    from_str='notify@digital.cabinet-office.gov.uk',
+                                    message=email_code,
+                                    subject='Verification code',
+                                    token=admin_api_client.auth_token)
+    except:
+        raise AdminApiClientException('Exception when sending email.')
+
+    return email_code
\ No newline at end of file
diff --git a/app/main/views/register.py b/app/main/views/register.py
index 953d514b5..c684e018c 100644
--- a/app/main/views/register.py
+++ b/app/main/views/register.py
@@ -1,15 +1,14 @@
 from datetime import datetime, timedelta
-from random import randint
 
 from flask import render_template, redirect, jsonify, session
 from sqlalchemy.exc import SQLAlchemyError
 
-from app import admin_api_client
 from app.main import main
 from app.main.dao import users_dao
 from app.main.encryption import hashpw
 from app.main.exceptions import AdminApiClientException
 from app.main.forms import RegisterUserForm
+from app.main.views import send_sms_code, send_email_code
 from app.models import User
 
 
@@ -46,28 +45,4 @@ def process_register():
     return redirect('/verify')
 
 
-def send_sms_code(mobile_number):
-    sms_code = _create_code()
-    try:
-        admin_api_client.send_sms(mobile_number, message=sms_code, token=admin_api_client.auth_token)
-    except:
-        raise AdminApiClientException('Exception when sending sms.')
-    return sms_code
 
-
-def send_email_code(email):
-    email_code = _create_code()
-    try:
-        admin_api_client.send_email(email_address=email,
-                                    from_str='notify@digital.cabinet-office.gov.uk',
-                                    message=email_code,
-                                    subject='Verification code',
-                                    token=admin_api_client.auth_token)
-    except:
-        raise AdminApiClientException('Exception when sending email.')
-
-    return email_code
-
-
-def _create_code():
-    return ''.join(["%s" % randint(0, 9) for _ in range(0, 5)])
diff --git a/app/main/views/sign_in.py b/app/main/views/sign_in.py
index aaa6d827b..8deebecd4 100644
--- a/app/main/views/sign_in.py
+++ b/app/main/views/sign_in.py
@@ -1,10 +1,12 @@
 from flask import render_template, redirect, jsonify
-from flask_login import login_user
+from flask import session
 
 from app.main import main
 from app.main.dao import users_dao
 from app.main.encryption import checkpw
+from app.main.encryption import hashpw
 from app.main.forms import LoginForm
+from app.main.views import send_sms_code
 
 
 @main.route("/sign-in", methods=(['GET']))
@@ -24,7 +26,9 @@ def process_sign_in():
         if not user.is_active():
             return jsonify(active_user=False), 401
         if checkpw(form.password.data, user.password):
-            login_user(user)
+            sms_code = send_sms_code(user.mobile_number)
+            session['user_id'] = user.id
+            session['sms_code'] = hashpw(sms_code)
         else:
             users_dao.increment_failed_login_count(user.id)
             return jsonify(authorization=False), 401
diff --git a/app/main/views/two_factor.py b/app/main/views/two_factor.py
index bac755107..f96de1661 100644
--- a/app/main/views/two_factor.py
+++ b/app/main/views/two_factor.py
@@ -1,4 +1,5 @@
 from flask import render_template, redirect, jsonify
+from flask_login import login_user
 
 from app.main import main
 from app.main.forms import TwoFactorForm
@@ -14,6 +15,7 @@ def process_two_factor():
     form = TwoFactorForm()
 
     if form.validate_on_submit():
+        login_user(user)
         return redirect('/dashboard')
     else:
         return jsonify(form.errors), 400
diff --git a/tests/app/main/views/test_two_factor.py b/tests/app/main/views/test_two_factor.py
index f7ea0c4a3..bc99d4c9e 100644
--- a/tests/app/main/views/test_two_factor.py
+++ b/tests/app/main/views/test_two_factor.py
@@ -8,7 +8,7 @@ def test_should_render_two_factor_page(notifications_admin, notifications_admin_
 
 def test_should_login_user_and_redirect_to_dashboard(notifications_admin, notifications_admin_db):
     response = notifications_admin.test_client().post('/two-factor',
-                                                       data={'sms_code': '12345'})
+                                                      data={'sms_code': '12345'})
 
     assert response.status_code == 302
-    assert response.location == 'http://localhost/dashboard'
\ No newline at end of file
+    assert response.location == 'http://localhost/dashboard'