MPT-26340 Add mfa caveat to macaroon token

auth/auth_server/auth_token/macaroon.py

@@ -3,6 +3,7 @@
 from pymacaroons import Macaroon, Verifier
 
 LOG = logging.getLogger(__name__)
+MFA_NONE = 'none'
 
 
 class MacaroonToken(object):
@@ -12,7 +13,7 @@ def __init__(self, secret, ident, location=''):
         self._ident = ident
         self._location = location
 
-    def create(self, register, provider):
+    def create(self, register, provider, mfa=MFA_NONE):
         macaroon = Macaroon(
             location=self._location,
             identifier=self._ident,
@@ -21,6 +22,7 @@ def create(self, register, provider):
         macaroon.add_first_party_caveat(f'created:{time.time()}')
         macaroon.add_first_party_caveat(f'register:{register}')
         macaroon.add_first_party_caveat(f'provider:{provider}')
+        macaroon.add_first_party_caveat(f'mfa:{mfa}')
         return macaroon.serialize()
 
     def verify(self, token):

auth/auth_server/controllers/token.py

@@ -8,7 +8,7 @@
 from auth.auth_server.controllers.base_async import BaseAsyncControllerWrapper
 from auth.auth_server.exceptions import Err
 from auth.auth_server.models.models import Token, User
-from auth.auth_server.auth_token.macaroon import MacaroonToken
+from auth.auth_server.auth_token.macaroon import MacaroonToken, MFA_NONE
 from auth.auth_server.utils import (hash_password, popkey,
                                     raise_not_provided_error)
 
@@ -102,13 +102,14 @@ def create_token_by_user_id(self, **kwargs):
             raise ForbiddenException(Err.OA0073, [])
         return self.create_user_token(user, **kwargs)
 
-    def create_user_token(self, user, **kwargs):
+    def create_user_token(self, user, mfa=MFA_NONE, **kwargs):
         model_type = self._get_model_type()
         LOG.info("Creating %s with parameters %s", model_type.__name__, kwargs)
         now = utcnow()
         macaroon_token = MacaroonToken(user.salt, user.id).create(
             xstr(kwargs.get('register', False)),
-            xstr(kwargs.get('provider', 'optscale'))
+            xstr(kwargs.get('provider', 'optscale')),
+            mfa,
         )
         params = {
             'user_id': user.id,

auth/auth_server/tests/unittests/test_token.py

@@ -1,5 +1,5 @@
 import uuid
-from auth.auth_server.auth_token.macaroon import MacaroonToken
+from auth.auth_server.auth_token.macaroon import MacaroonToken, MFA_NONE
 from auth.auth_server.models.models import Type, User
 from auth.auth_server.tests.unittests.test_model_base import TestModelBase
 from auth.auth_server.tests.unittests.utils import extract_caveats
@@ -41,6 +41,7 @@ def _check_token(self, user):
         caveats = extract_caveats(token)
         self.assertEqual(provider, caveats['provider'])
         self.assertEqual(register, caveats['register'])
+        self.assertEqual(MFA_NONE, caveats['mfa'])
 
     def test_token_root(self):
         user_root = User('root@hystax.com', self.type_root, 'pass',
@@ -86,3 +87,15 @@ def test_token_forgery(self):
         # pylint: disable=W0212
         macaroon_token._secret = 'fake'
         self.assertFalse(macaroon_token.verify(token))
+
+    def test_token_mfa_caveat(self):
+        session = self.db_session
+        user = User('mfa@hystax.com', self.type_customer, 'pass',
+                    display_name='MFA user', scope_id=str(uuid.uuid4()))
+        session.add(user)
+        session.commit()
+        macaroon_token = MacaroonToken(user.salt, user.id)
+        token = macaroon_token.create(register=False, provider='optscale',
+                                      mfa='totp')
+        self.assertTrue(macaroon_token.verify(token))
+        self.assertEqual('totp', extract_caveats(token)['mfa'])