From 9121081f4f7dc9b58aa5d818e32a14b02cb26f33 Mon Sep 17 00:00:00 2001
From: Michiel <michiel@deroos.ca>
Date: Wed, 19 May 2021 15:23:34 -0300
Subject: [PATCH] Update TFA Logic (#698)

* Update TFA Logic

* version

* comment

* Test
---
 app/server/models/user.py                    |  1 +
 app/server/utils/auth.py                     | 10 ++++++++++
 app/test_app/functional/api/test_auth_api.py | 12 ++++++++++++
 3 files changed, 23 insertions(+)

diff --git a/app/server/models/user.py b/app/server/models/user.py
index 2738cc0cd..417ce8eae 100644
--- a/app/server/models/user.py
+++ b/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
diff --git a/app/server/utils/auth.py b/app/server/utils/auth.py
index e38aa6562..cb234c49a 100644
--- a/app/server/utils/auth.py
+++ b/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 = {
diff --git a/app/test_app/functional/api/test_auth_api.py b/app/test_app/functional/api/test_auth_api.py
index fec1e0d13..c817477ab 100644
--- a/app/test_app/functional/api/test_auth_api.py
+++ b/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):
     """