109526036: Persist the verify code to the db.

The codes are hashed and saved to the db.
The code is marked as used once a valid code is submitted.
The code is valid for 1 hour.
The codes are no longer saved to the session.

app/main/dao/verify_codes_dao.py

@@ -0,0 +1,41 @@
+from datetime import datetime, timedelta
+
+from app import db
+from app.main.encryption import hashpw
+from app.models import VerifyCodes
+
+
+def add_code(user_id, code, code_type):
+    code = VerifyCodes(user_id=user_id,
+                       code=hashpw(code),
+                       code_type=code_type,
+                       expiry_datetime=datetime.now() + timedelta(hours=1))
+
+    db.session.add(code)
+    db.session.commit()
+
+
+def get_code(user_id, code_type):
+    verify_code = VerifyCodes.query.filter_by(user_id=user_id, code_type=code_type, code_used=False).first()
+    return verify_code
+
+
+def get_code_by_code(user_id, code_type):
+    return VerifyCodes.query.filter_by(user_id=user_id, code_type=code_type).first()
+
+
+def use_code(id):
+    verify_code = VerifyCodes.query.filter_by(id=id).first()
+    verify_code.code_used = True
+    db.session.add(verify_code)
+    db.session.commit()
+
+
+def add_code_with_expiry(user_id, code, code_type, expiry):
+    code = VerifyCodes(user_id=user_id,
+                       code=code,
+                       code_type=code_type,
+                       expiry_datetime=expiry)
+
+    db.session.add(code)
+    db.session.commit()

app/main/forms.py

@@ -1,8 +1,11 @@
+from datetime import datetime
+
 from flask import session
 from flask_wtf import Form
 from wtforms import StringField, PasswordField
 from wtforms.validators import DataRequired, Email, Length, Regexp
 
+from app.main.dao import verify_codes_dao
 from app.main.encryption import checkpw
 from app.main.validators import Blacklist
 
@@ -46,7 +49,8 @@ class TwoFactorForm(Form):
                                                    Regexp(regex=verify_code, message='Code must be 5 digits')])
 
     def validate_sms_code(self, a):
-        validate_code(self.sms_code, session['sms_code'])
+        code = verify_codes_dao.get_code(session['user_id'], 'sms')
+        validate_code(self.sms_code, code)
 
 
 class VerifyForm(Form):
@@ -58,18 +62,24 @@ class VerifyForm(Form):
                                          Regexp(regex=verify_code, message='Code must be 5 digits')])
 
     def validate_email_code(self, a):
-        validate_code(self.email_code, session['email_code'])
+        code = verify_codes_dao.get_code(session['user_id'], 'email')
+        validate_code(self.email_code, code)
 
     def validate_sms_code(self, a):
-        validate_code(self.sms_code, session['sms_code'])
+        code = verify_codes_dao.get_code(session['user_id'], 'sms')
+        validate_code(self.sms_code, code)
 
 
 def validate_code(field, code):
+    if code.expiry_datetime < datetime.now():
+        field.errors.append('Code has expired')
+        return False
     if field.data is not None:
-        if checkpw(str(field.data), code) is False:
+        if checkpw(field.data, code.code) is False:
             field.errors.append('Code does not match')
             return False
         else:
+            verify_codes_dao.use_code(code.id)
             return True
     else:
         return True

app/main/views/__init__.py

@@ -1,24 +1,27 @@
 from random import randint
 from app import admin_api_client
 from app.main.exceptions import AdminApiClientException
+from app.main.dao import verify_codes_dao
 
 
 def create_verify_code():
     return ''.join(["%s" % randint(0, 9) for _ in range(0, 5)])
 
 
-def send_sms_code(mobile_number):
+def send_sms_code(user_id, mobile_number):
     sms_code = create_verify_code()
     try:
+        verify_codes_dao.add_code(user_id=user_id, code=sms_code, code_type='sms')
         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):
+def send_email_code(user_id, email):
     email_code = create_verify_code()
     try:
+        verify_codes_dao.add_code(user_id=user_id, code=email_code, code_type='email')
         admin_api_client.send_email(email_address=email,
                                     from_str='notify@digital.cabinet-office.gov.uk',
                                     message=email_code,

app/main/views/register.py

@@ -29,12 +29,10 @@ def process_register():
                     created_at=datetime.now(),
                     role_id=1)
         try:
-            sms_code = send_sms_code(form.mobile_number.data)
-            email_code = send_email_code(form.email_address.data)
-            session['sms_code'] = hashpw(sms_code)
-            session['email_code'] = hashpw(email_code)
-            session['expiry_date'] = str(datetime.now() + timedelta(hours=1))
             users_dao.insert_user(user)
+            send_sms_code(user_id=user.id, mobile_number=form.mobile_number.data)
+            send_email_code(user_id=user.id, email=form.email_address.data)
+            session['expiry_date'] = str(datetime.now() + timedelta(hours=1))
             session['user_id'] = user.id
         except AdminApiClientException as e:
             return jsonify(admin_api_client_error=e.value)

app/main/views/sign_in.py

@@ -26,7 +26,7 @@ def process_sign_in():
         if not user.is_active():
             return jsonify(active_user=False), 401
         if checkpw(form.password.data, user.password):
-            sms_code = send_sms_code(user.mobile_number)
+            sms_code = send_sms_code(user.id, user.mobile_number)
             session['user_id'] = user.id
             session['sms_code'] = hashpw(sms_code)
         else:

app/models.py

@@ -5,6 +5,19 @@ DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 DATE_FORMAT = "%Y-%m-%d"
 
 
+class VerifyCodes(db.Model):
+    __tablename__ = 'verify_codes'
+
+    code_types = ['email', 'sms']
+
+    id = db.Column(db.Integer, primary_key=True)
+    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), index=True, unique=False, nullable=False)
+    code = db.Column(db.String, nullable=False)
+    code_type = db.Column(db.Enum(code_types, name='verify_code_types'), index=False, unique=False, nullable=False)
+    expiry_datetime = db.Column(db.DateTime, nullable=False)
+    code_used = db.Column(db.Boolean, default=False)
+
+
 class Roles(db.Model):
     __tablename__ = 'roles'