diff --git a/auth/auth_server/auth_token/macaroon.py b/auth/auth_server/auth_token/macaroon.py index f1cd23222..ebec9518c 100644 --- a/auth/auth_server/auth_token/macaroon.py +++ b/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): diff --git a/auth/auth_server/controllers/token.py b/auth/auth_server/controllers/token.py index 455a1daef..825de1ad8 100644 --- a/auth/auth_server/controllers/token.py +++ b/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, diff --git a/auth/auth_server/tests/unittests/test_token.py b/auth/auth_server/tests/unittests/test_token.py index c16e5003a..07af73ca0 100644 --- a/auth/auth_server/tests/unittests/test_token.py +++ b/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'])