Merge pull request #1725 from TimoGlastra/feat/rotate-wallet-token

feat: create new JWT tokens and invalidate older for multitenancy

Multitenancy.md

@@ -260,7 +260,7 @@ echo $new_tenant | curl -X POST "${ACAPY_ADMIN_URL}/multitenancy/wallet" \
 
 #### Method 2: Get tenant token
 
-This method allows you to retrieve a tenant `token` for an already registered tenant.  To retrieve a token you will need an Admin API key (if your admin is protected with one), `wallet_key` and the `wallet_id` of the tenant.
+This method allows you to retrieve a tenant `token` for an already registered tenant.  To retrieve a token you will need an Admin API key (if your admin is protected with one), `wallet_key` and the `wallet_id` of the tenant. Note that calling the get tenant token endpoint will **invalidate** the old token. This is useful if the old token needs to be revoked, but does mean that you can't have multiple authentication tokens for the same wallet. Only the last generated token will always be valid.
 
 Example
 

aries_cloudagent/multitenant/admin/routes.py

@@ -315,7 +315,7 @@ async def wallet_create(request: web.BaseRequest):
             settings, key_management_mode
         )
 
-        token = multitenant_mgr.create_auth_token(wallet_record, wallet_key)
+        token = await multitenant_mgr.create_auth_token(wallet_record, wallet_key)
     except BaseError as err:
         raise web.HTTPBadRequest(reason=err.roll_up) from err
 
@@ -413,7 +413,7 @@ async def wallet_create_token(request: web.BaseRequest):
                 " the wallet key to be provided"
             )
 
-        token = multitenant_mgr.create_auth_token(wallet_record, wallet_key)
+        token = await multitenant_mgr.create_auth_token(wallet_record, wallet_key)
     except StorageNotFoundError as err:
         raise web.HTTPNotFound(reason=err.roll_up) from err
     except WalletKeyMissingError as err:

aries_cloudagent/multitenant/admin/tests/test_routes.py

@@ -164,7 +164,7 @@ async def test_wallet_create(self):
                 return_value=wallet_mock
             )
 
-            self.mock_multitenant_mgr.create_auth_token = async_mock.Mock(
+            self.mock_multitenant_mgr.create_auth_token = async_mock.CoroutineMock(
                 return_value="test_token"
             )
             print(self.request["context"])
@@ -215,7 +215,7 @@ async def test_wallet_create_optional_default_fields(self):
 
         with async_mock.patch.object(test_module.web, "json_response") as mock_response:
             self.mock_multitenant_mgr.create_wallet = async_mock.CoroutineMock()
-            self.mock_multitenant_mgr.create_auth_token = async_mock.Mock()
+            self.mock_multitenant_mgr.create_auth_token = async_mock.CoroutineMock()
 
             await test_module.wallet_create(self.request)
             self.mock_multitenant_mgr.create_wallet.assert_called_once_with(
@@ -440,7 +440,7 @@ async def test_wallet_create_token_managed(self):
         ) as mock_response:
             mock_wallet_record_retrieve_by_id.return_value = mock_wallet_record
 
-            self.mock_multitenant_mgr.create_auth_token = async_mock.Mock(
+            self.mock_multitenant_mgr.create_auth_token = async_mock.CoroutineMock(
                 return_value="test_token"
             )
 
@@ -468,7 +468,7 @@ async def test_wallet_create_token_unmanaged(self):
         ) as mock_response:
             mock_wallet_record_retrieve_by_id.return_value = mock_wallet_record
 
-            self.mock_multitenant_mgr.create_auth_token = async_mock.Mock(
+            self.mock_multitenant_mgr.create_auth_token = async_mock.CoroutineMock(
                 return_value="test_token"
             )
 

aries_cloudagent/multitenant/base.py

@@ -1,5 +1,6 @@
 """Manager for multitenancy."""
 
+from datetime import datetime
 import logging
 from abc import abstractmethod
 
@@ -333,7 +334,7 @@ async def add_key(
                 keylist_updates, connection_id=mediation_record.connection_id
             )
 
-    def create_auth_token(
+    async def create_auth_token(
         self, wallet_record: WalletRecord, wallet_key: str = None
     ) -> str:
         """Create JWT auth token for specified wallet record.
@@ -351,8 +352,9 @@ def create_auth_token(
             str: JWT auth token
 
         """
+        iat = int(round(datetime.utcnow().timestamp()))
 
-        jwt_payload = {"wallet_id": wallet_record.wallet_id}
+        jwt_payload = {"wallet_id": wallet_record.wallet_id, "iat": iat}
         jwt_secret = self._profile.settings.get("multitenant.jwt_secret")
 
         if wallet_record.requires_external_key:
@@ -363,6 +365,11 @@ def create_auth_token(
 
         token = jwt.encode(jwt_payload, jwt_secret, algorithm="HS256").decode()
 
+        # Store iat for verification later on
+        wallet_record.jwt_iat = iat
+        async with self._profile.session() as session:
+            await wallet_record.save(session)
+
         return token
 
     async def get_profile_for_token(
@@ -389,6 +396,7 @@ async def get_profile_for_token(
 
         wallet_id = token_body.get("wallet_id")
         wallet_key = token_body.get("wallet_key")
+        iat = token_body.get("iat")
 
         async with self._profile.session() as session:
             wallet = await WalletRecord.retrieve_by_id(session, wallet_id)
@@ -399,6 +407,9 @@ async def get_profile_for_token(
 
             extra_settings["wallet.key"] = wallet_key
 
+        if wallet.jwt_iat and wallet.jwt_iat != iat:
+            raise MultitenantManagerError("Token not valid")
+
         profile = await self.get_wallet_profile(context, wallet, extra_settings)
 
         return profile

aries_cloudagent/multitenant/tests/test_base.py

@@ -3,6 +3,8 @@
 
 import jwt
 
+from datetime import datetime
+
 from ...core.in_memory import InMemoryProfile
 from ...config.base import InjectionError
 from ...messaging.responder import BaseResponder
@@ -21,6 +23,7 @@
 from ...wallet.did_method import DIDMethod
 from ..base import BaseMultitenantManager, MultitenantManagerError
 from ..error import WalletKeyMissingError
+from .. import base as test_module
 
 
 class TestBaseMultitenantManager(AsyncTestCase):
@@ -414,35 +417,55 @@ async def test_create_auth_token_fails_no_wallet_key_but_required(self):
 
     async def test_create_auth_token_managed(self):
         self.profile.settings["multitenant.jwt_secret"] = "very_secret_jwt"
-        wallet_record = WalletRecord(
+        wallet_record = async_mock.MagicMock(
             wallet_id="test_wallet",
             key_management_mode=WalletRecord.MODE_MANAGED,
+            requires_external_key=False,
             settings={},
+            save=async_mock.CoroutineMock(),
         )
 
+        utc_now = datetime(2020, 1, 1, 0, 0, 0)
+        iat = int(round(utc_now.timestamp()))
+
         expected_token = jwt.encode(
-            {"wallet_id": wallet_record.wallet_id}, "very_secret_jwt"
+            {"wallet_id": wallet_record.wallet_id, "iat": iat}, "very_secret_jwt"
         ).decode()
 
-        token = self.manager.create_auth_token(wallet_record)
+        with async_mock.patch.object(test_module, "datetime") as mock_datetime:
+            mock_datetime.utcnow.return_value = utc_now
+            token = await self.manager.create_auth_token(wallet_record)
 
+        assert wallet_record.jwt_iat == iat
         assert expected_token == token
 
     async def test_create_auth_token_unmanaged(self):
         self.profile.settings["multitenant.jwt_secret"] = "very_secret_jwt"
-        wallet_record = WalletRecord(
+        wallet_record = async_mock.MagicMock(
             wallet_id="test_wallet",
             key_management_mode=WalletRecord.MODE_UNMANAGED,
+            requires_external_key=True,
             settings={"wallet.type": "indy"},
+            save=async_mock.CoroutineMock(),
         )
 
+        utc_now = datetime(2020, 1, 1, 0, 0, 0)
+        iat = int(round(utc_now.timestamp()))
+
         expected_token = jwt.encode(
-            {"wallet_id": wallet_record.wallet_id, "wallet_key": "test_key"},
+            {
+                "wallet_id": wallet_record.wallet_id,
+                "iat": iat,
+                "wallet_key": "test_key",
+            },
             "very_secret_jwt",
         ).decode()
 
-        token = self.manager.create_auth_token(wallet_record, "test_key")
+        with async_mock.patch.object(test_module, "datetime") as mock_datetime:
+            mock_datetime.utcnow.return_value = utc_now
+            token = await self.manager.create_auth_token(wallet_record, "test_key")
 
+        assert wallet_record.jwt_iat == iat
         assert expected_token == token
 
     async def test_get_profile_for_token_invalid_token_raises(self):
@@ -468,7 +491,7 @@ async def test_get_profile_for_token_wallet_key_missing_raises(self):
         with self.assertRaises(WalletKeyMissingError):
             await self.manager.get_profile_for_token(self.profile.context, token)
 
-    async def test_get_profile_for_token_managed_wallet(self):
+    async def test_get_profile_for_token_managed_wallet_no_iat(self):
         self.profile.settings["multitenant.jwt_secret"] = "very_secret_jwt"
         wallet_record = WalletRecord(
             key_management_mode=WalletRecord.MODE_MANAGED,
@@ -500,6 +523,83 @@ async def test_get_profile_for_token_managed_wallet(self):
 
             assert profile == mock_profile
 
+    async def test_get_profile_for_token_managed_wallet_iat(self):
+        iat = 100
+
+        self.profile.settings["multitenant.jwt_secret"] = "very_secret_jwt"
+        wallet_record = WalletRecord(
+            key_management_mode=WalletRecord.MODE_MANAGED,
+            settings={"wallet.type": "indy", "wallet.key": "wallet_key"},
+            jwt_iat=iat,
+        )
+
+        session = await self.profile.session()
+        await wallet_record.save(session)
+
+        token = jwt.encode(
+            {"wallet_id": wallet_record.wallet_id, "iat": iat},
+            "very_secret_jwt",
+            algorithm="HS256",
+        ).decode()
+
+        with async_mock.patch.object(
+            BaseMultitenantManager, "get_wallet_profile"
+        ) as get_wallet_profile:
+            mock_profile = InMemoryProfile.test_profile()
+            get_wallet_profile.return_value = mock_profile
+
+            profile = await self.manager.get_profile_for_token(
+                self.profile.context, token
+            )
+
+            get_wallet_profile.assert_called_once_with(
+                self.profile.context,
+                wallet_record,
+                {},
+            )
+
+            assert profile == mock_profile
+
+    async def test_get_profile_for_token_managed_wallet_x_iat_no_match(self):
+        iat = 100
+
+        self.profile.settings["multitenant.jwt_secret"] = "very_secret_jwt"
+        wallet_record = WalletRecord(
+            key_management_mode=WalletRecord.MODE_MANAGED,
+            settings={"wallet.type": "indy", "wallet.key": "wallet_key"},
+            jwt_iat=iat,
+        )
+
+        session = await self.profile.session()
+        await wallet_record.save(session)
+
+        token = jwt.encode(
+            # Change iat from record value
+            {"wallet_id": wallet_record.wallet_id, "iat": 200},
+            "very_secret_jwt",
+            algorithm="HS256",
+        ).decode()
+
+        with async_mock.patch.object(
+            BaseMultitenantManager, "get_wallet_profile"
+        ) as get_wallet_profile, self.assertRaises(
+            MultitenantManagerError, msg="Token not valid"
+        ):
+            mock_profile = InMemoryProfile.test_profile()
+            get_wallet_profile.return_value = mock_profile
+
+            profile = await self.manager.get_profile_for_token(
+                self.profile.context, token
+            )
+
+            get_wallet_profile.assert_called_once_with(
+                self.profile.context,
+                wallet_record,
+                {},
+            )
+
+            assert profile == mock_profile
+
     async def test_get_profile_for_token_unmanaged_wallet(self):
         self.profile.settings["multitenant.jwt_secret"] = "very_secret_jwt"
         wallet_record = WalletRecord(

aries_cloudagent/wallet/models/wallet_record.py

@@ -39,11 +39,13 @@ def __init__(
         # MTODO: how to make this a tag without making it
         # a constructor param
         wallet_name: str = None,
+        jwt_iat: Optional[int] = None,
         **kwargs,
     ):
         """Initialize a new WalletRecord."""
         super().__init__(wallet_id, **kwargs)
         self.key_management_mode = key_management_mode
+        self.jwt_iat = jwt_iat
         self._settings = settings
 
     @property
@@ -85,7 +87,8 @@ def wallet_key(self) -> Optional[str]:
     def record_value(self) -> dict:
         """Accessor for the JSON record value generated for this record."""
         return {
-            prop: getattr(self, prop) for prop in ("settings", "key_management_mode")
+            prop: getattr(self, prop)
+            for prop in ("settings", "key_management_mode", "jwt_iat")
         }
 
     @property