Update TFA Logic (#698)

* Update TFA Logic

* version

* comment

* Test

app/server/models/user.py

@@ -497,6 +497,7 @@ def encode_TFA_token(self, valid_days=1):
         try:
 
             payload = {
+                'token_type': 'TFA',
                 'exp': datetime.datetime.utcnow() + datetime.timedelta(days=valid_days, seconds=30),
                 'iat': datetime.datetime.utcnow(),
                 'id': self.id

app/server/utils/auth.py

@@ -279,6 +279,7 @@ def tfa_logic(user, tfa_token, ignore_tfa_requirement=False):
             return response_object
 
         tfa_response = User.decode_auth_token(tfa_token, 'TFA')
+
         if isinstance(tfa_response, str):
             # User doesn't have valid TFA token
             response_object = {
@@ -287,6 +288,15 @@ def tfa_logic(user, tfa_token, ignore_tfa_requirement=False):
             }
             return response_object
 
+        if tfa_response.get("token_type") != "TFA":
+            # Ensure the TFA token was generated by encode_TFA_token
+            response_object = {
+                'tfa_failure': True,
+                'message': 'Invalid TFA response'
+            }
+
+            return response_object
+
         if tfa_response.get("id") != user.id:
             # User doesn't has valid TFA token BUT it's not theirs
             response_object = {

app/test_app/functional/api/test_auth_api.py

@@ -367,6 +367,18 @@ def test_get_external_credentials_api(test_client, authed_sempo_admin_user):
     org = Organisation.query.filter_by(external_auth_username = response.json['username']).first()
     assert response.json['password'] == org.external_auth_password
 
+def test_tfa_token_integrity(test_client, authed_sempo_admin_user):
+    """
+    Ensure that when a phony TFA token is provided, an error is returned
+    """
+    authed_sempo_admin_user.is_activated = True
+    authed_sempo_admin_user.TFA_enabled = True
+    authed_sempo_admin_user.set_held_role('ADMIN', 'admin')
+    auth_token = authed_sempo_admin_user.encode_auth_token().decode()
+    response = test_client.get('/api/v1/auth/external/',
+                               headers=dict(Authorization=auth_token + '|' + auth_token, Accept='application/json'),
+                               content_type='application/json', follow_redirects=True)
+    assert response.json['message'] == 'Invalid TFA response'
 
 def test_logout_api(test_client, authed_sempo_admin_user):
     """