Repository: dwavesystems/dwave-cloud-client
Pull request #609 (2024-03-05 22:17:04)

Implement OAuth 2.0 token revocation
Close #600.

Note: token revocation is currently only implemented in the EXP environment.


Patch:
diff --git a/dwave/cloud/auth/flows.py b/dwave/cloud/auth/flows.py
index 7403aea..be2e20d 100644
--- a/dwave/cloud/auth/flows.py
+++ b/dwave/cloud/auth/flows.py
@@ -18,10 +18,11 @@ import logging
 import time
 import webbrowser
 from operator import sub
-from typing import Any, Callable, Dict, Optional, Union, Sequence
+from typing import Any, Callable, Dict, Optional, Union, Sequence, Literal
 from urllib.parse import urljoin
 
 import click
+import requests
 from authlib.integrations.requests_client import OAuth2Session, OAuthError
 from authlib.oauth2.rfc6749 import OAuth2Token
 from authlib.common.security import generate_token
@@ -55,6 +56,8 @@ class AuthFlow:
             URL of the authorization server's authorization endpoint.
         token_endpoint:
             URL of the authorization server's token endpoint.
+        revocation_endpoint:
+            URL of the authorization server's OAuth 2.0 revocation endpoint.
         session_config:
             Configuration options for the low-level ``requests.Session`` used
             for all OAuth 2 requests. Supported options are: ``cert``,
@@ -78,6 +81,7 @@ class AuthFlow:
                  redirect_uri: str,
                  authorization_endpoint: str,
                  token_endpoint: str,
+                 revocation_endpoint: Optional[str] = None,
                  session_config: Optional[Dict[str, Any]] = None,
                  leap_api_endpoint: Optional[str] = None,
                  creds: Optional[Credentials] = None
@@ -86,9 +90,12 @@ class AuthFlow:
         self.scopes = ' '.join(scopes)
         self.authorization_endpoint = authorization_endpoint
         self.token_endpoint = token_endpoint
+        self.revocation_endpoint = revocation_endpoint
         self.leap_api_endpoint = leap_api_endpoint
         self.creds = creds
 
+        # TODO: verify authorization/token/revocation endpoints are https
+
         self.session = OAuth2Session(
             client_id=client_id, scope=scopes, redirect_uri=redirect_uri,
             code_challenge_method='S256',
@@ -205,6 +212,39 @@ class AuthFlow:
         self._save_token_to_creds(token)
         return token
 
+    def revoke_token(self, *,
+                     token: Optional[str] = None,
+                     token_type_hint: Optional[
+                         Literal['access_token', 'refresh_token']] = None
+                     ) -> bool:
+        """Revoke OAuth 2.0 token.
+
+        Args:
+            token:
+                OAuth 2.0 access or refresh token to be revoked. If refresh token
+                is specified, the authorization server should also revoke the
+                corresponding access tokens. If ``token`` is not specified, then
+                current refresh token is revoked.
+            token_type_hint:
+                Type of the token to be revoked.
+
+        Returns:
+            Token revocation status. True if revocation was successful, or token
+            was invalid already. False in case of an error.
+
+        Note:
+            Token revocation requires that a ``revocation_endpoint`` is defined
+            in :class:`~dwave.cloud.auth.flows.AuthFlow` construction.
+        """
+        if self.revocation_endpoint is None:
+            raise TypeError("Revocation endpoint undefined.")
+
+        response: requests.Response = self.session.revoke_token(
+            url=self.revocation_endpoint, token=token,
+            token_type_hint=token_type_hint)
+
+        return response.ok
+
     def ensure_active_token(self):
         if not self.token:
             return False
@@ -259,6 +299,10 @@ class LeapAuthFlow(AuthFlow):
     def _infer_token_endpoint(leap_api_endpoint: str) -> str:
         return urljoin(leap_api_endpoint, '/leap/openid/token')
 
+    @staticmethod
+    def _infer_revocation_endpoint(leap_api_endpoint: str) -> str:
+        return urljoin(leap_api_endpoint, '/leap/openid/revoke_token/')
+
     @staticmethod
     def _infer_leap_success_uri(leap_api_endpoint: str) -> str:
         return urljoin(leap_api_endpoint, '/leap/openid/success/')
@@ -290,6 +334,7 @@ class LeapAuthFlow(AuthFlow):
 
         authorization_endpoint = cls._infer_auth_endpoint(config.leap_api_endpoint)
         token_endpoint = cls._infer_token_endpoint(config.leap_api_endpoint)
+        revocation_endpoint = cls._infer_revocation_endpoint(config.leap_api_endpoint)
 
         # note: possible to de-dup via `DWaveAPIClient.from_config_model`, but
         # currently not worth the extra complexity (indirection layer).
@@ -306,6 +351,7 @@ class LeapAuthFlow(AuthFlow):
             redirect_uri=kwargs.pop('redirect_uri', cls._OOB_REDIRECT_URI),
             authorization_endpoint=authorization_endpoint,
             token_endpoint=token_endpoint,
+            revocation_endpoint=revocation_endpoint,
             session_config=session_config,
             leap_api_endpoint=config.leap_api_endpoint,
             creds=creds)
diff --git a/dwave/cloud/cli.py b/dwave/cloud/cli.py
index 08bfb44..5478cf7 100644
--- a/dwave/cloud/cli.py
+++ b/dwave/cloud/cli.py
@@ -1009,11 +1009,11 @@ def get(*, config_file, profile, token_type, json_output, output):
     flow = LeapAuthFlow.from_config_model(config)
 
     token_key = token_type.replace('-', '_')
-    if not flow.token or not (token_value := flow.token.get(token_key)):
+    if not flow.token or not token_key in flow.token:
         raise CLIError('Token not found. Please run "dwave auth login".', code=100)
 
     token_pretty = token_type.replace('-', ' ').capitalize()
-    output(f"{token_pretty}: {{%s}}" % token_key , **{token_key: token_value})
+    output(f"{token_pretty}: {{%s}}" % token_key, **{token_key: flow.token[token_key]})
 
     if token_type == 'access-token':
         expires_at = flow.token.get('expires_at')
@@ -1051,6 +1051,32 @@ def refresh(*, config_file, profile, output):
            'You can now use "dwave auth get" to view them.')
 
 
+@auth.command()
+@config_file_options()
+@click.argument('token_type', default='access-token',
+                type=click.Choice(['access-token', 'refresh-token']))
+@json_output
+@standardized_output
+def revoke(*, config_file, profile, token_type, json_output, output):
+    """Revoke access and/or refresh token(s) for Leap API."""
+
+    config = validate_config_v1(load_config(config_file=config_file, profile=profile))
+    flow = LeapAuthFlow.from_config_model(config)
+
+    token_key = token_type.replace('-', '_')
+    token_pretty = token_type.replace('-', ' ').capitalize()
+    if not flow.token or not token_key in flow.token:
+        raise CLIError('Token not found. Please run "dwave auth login".', code=100)
+
+    # revoke
+    revoked = flow.revoke_token(token=flow.token[token_key], token_type_hint=token_key)
+
+    if not revoked:
+        raise CLIError(f'{token_pretty} revocation failed.', code=102)
+
+    output(f'{token_pretty} successfully revoked.')
+
+
 @cli.group()
 def leap():
     """Interact with Leap API."""
diff --git a/releasenotes/notes/add-auth-token-revoke-080639f74044f49f.yaml b/releasenotes/notes/add-auth-token-revoke-080639f74044f49f.yaml
new file mode 100644
index 0000000..a92055a
--- /dev/null
+++ b/releasenotes/notes/add-auth-token-revoke-080639f74044f49f.yaml
@@ -0,0 +1,10 @@
+---
+features:
+  - |
+    Add ``revoke_token`` method to ``dwave.cloud.auth.flow.AuthFlow`` to enable
+    revocation of OAuth 2.0 refresh/access tokens.
+    See `#600 <https://github.com/dwavesystems/dwave-cloud-client/issues/600>`_.
+  - |
+    Add ``dwave auth revoke`` CLI command to enable revocation of OAuth 2.0
+    refresh/access tokens from command line.
+    See `#600 <https://github.com/dwavesystems/dwave-cloud-client/issues/600>`_.


Test patch:
diff --git a/tests/auth/test_flows.py b/tests/auth/test_flows.py
index 6031d96..ee46f50 100644
--- a/tests/auth/test_flows.py
+++ b/tests/auth/test_flows.py
@@ -38,6 +38,7 @@ class TestAuthFlow(unittest.TestCase):
         self.redirect_uri_oob = 'oob'
         self.authorization_endpoint = 'https://example.com/authorize'
         self.token_endpoint = 'https://example.com/token'
+        self.revocation_endpoint = 'https://example.com/revoke'
         self.token = dict(access_token='123', refresh_token='456', id_token='789')
         self.creds = Credentials(create=False)
         self.leap_api_endpoint = 'https://example.com/leap/api'
@@ -48,6 +49,7 @@ class TestAuthFlow(unittest.TestCase):
             redirect_uri=self.redirect_uri_oob,
             authorization_endpoint=self.authorization_endpoint,
             token_endpoint=self.token_endpoint,
+            revocation_endpoint=self.revocation_endpoint,
             leap_api_endpoint=self.leap_api_endpoint,
             creds=self.creds)
 
@@ -149,6 +151,89 @@ class TestAuthFlow(unittest.TestCase):
 
         m.assert_called_once_with(token=self.token)
 
+    def test_revoke_token(self):
+        flow = AuthFlow(**self.test_args)
+
+        ok = requests.Response()
+        ok.status_code = 200
+
+        with mock.patch.object(flow.session, 'revoke_token',
+                               return_value=ok) as m:
+            status = flow.revoke_token()
+            self.assertTrue(status)
+
+        m.assert_called_once_with(url=flow.revocation_endpoint, token=None,
+                                  token_type_hint=None)
+
+    def test_revoke_token_is_conditionally_available(self):
+        # when: `revocation_endpoint` is undefined
+        args = self.test_args.copy()
+        del args['revocation_endpoint']
+
+        # then: init works
+        flow = AuthFlow(**args)
+
+        # but: revoke_token() is unavailable
+        with self.assertRaisesRegex(TypeError, 'endpoint undefined'):
+            flow.revoke_token()
+
+    @requests_mock.Mocker()
+    def test_revoke_token_request_formation(self, m):
+        # when we want to revoke a specific token, outgoing request conforms to RFC 7009
+
+        # mock the revocation_endpoint
+        expected_params = dict(
+            token=self.token['access_token'],
+            token_type_hint='access_token')
+
+        def post_body_matcher(request):
+            params = dict(parse_qsl(request.text))
+            return params == expected_params
+
+        m.get(requests_mock.ANY, status_code=404)
+        m.post(requests_mock.ANY, status_code=404)
+        m.post(self.revocation_endpoint, json=dict(error="error", error_description="bad request"))
+        m.post(self.revocation_endpoint, additional_matcher=post_body_matcher, status_code=200)
+
+        # reset creds
+        self.creds.clear()
+
+        # initialize flow with token
+        flow = AuthFlow(**self.test_args)
+        flow.token = self.token
+
+        # verify revoke token flow
+        status = flow.revoke_token(token=self.token['access_token'],
+                                   token_type_hint='access_token')
+        self.assertTrue(status)
+
+    @requests_mock.Mocker()
+    def test_revoke_token_default_request_formation(self, m):
+        # when token is unspecified, refresh token should be revoked
+
+        # mock the revocation_endpoint:
+        expected_params = dict(token=self.token['refresh_token'])
+
+        def post_body_matcher(request):
+            params = dict(parse_qsl(request.text))
+            return params == expected_params
+
+        m.get(requests_mock.ANY, status_code=404)
+        m.post(requests_mock.ANY, status_code=404)
+        m.post(self.revocation_endpoint, json=dict(error="error", error_description="bad request"))
+        m.post(self.revocation_endpoint, additional_matcher=post_body_matcher, status_code=200)
+
+        # reset creds
+        self.creds.clear()
+
+        # initialize flow with token
+        flow = AuthFlow(**self.test_args)
+        flow.token = self.token
+
+        # verify revoke token flow
+        status = flow.revoke_token()
+        self.assertTrue(status)
+
     def test_session_config(self):
         config = dict(
             cert='/path/to/cert',
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 87d0b5b..23aede1 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -515,6 +515,8 @@ class TestAuthCli(unittest.TestCase):
     @mock.patch('dwave.cloud.auth.flows.LeapAuthFlow.from_config_model')
     def test_get(self, flow_factory):
         flow = flow_factory.return_value
+        token = dict(access_token='123', refresh_token='456')
+        type(flow).token = mock.PropertyMock(return_value=token)
 
         with self.subTest('dwave auth get access-token'):
             flow.reset_mock()
@@ -522,11 +524,8 @@ class TestAuthCli(unittest.TestCase):
             with runner.isolated_filesystem():
                 result = runner.invoke(cli, ['auth', 'get', 'access-token'])
 
-            flow.token.get.assert_has_calls(
-                [mock.call('access_token'), mock.call('expires_at')], any_order=True)
-            self.assertIn('Access token:', result.output)
-
             self.assertEqual(result.exit_code, 0)
+            self.assertIn(token['access_token'], result.output)
 
         with self.subTest('dwave auth get refresh-token'):
             flow.reset_mock()
@@ -534,10 +533,8 @@ class TestAuthCli(unittest.TestCase):
             with runner.isolated_filesystem():
                 result = runner.invoke(cli, ['auth', 'get', 'refresh-token'])
 
-            flow.token.get.assert_called_with('refresh_token')
-            self.assertIn('Refresh token:', result.output)
-
             self.assertEqual(result.exit_code, 0)
+            self.assertIn(token['refresh_token'], result.output)
 
     @mock.patch('dwave.cloud.auth.flows.LeapAuthFlow.from_config_model')
     def test_refresh(self, flow_factory):
@@ -562,6 +559,62 @@ class TestAuthCli(unittest.TestCase):
 
             self.assertEqual(result.exit_code, 100)
 
+    @mock.patch('dwave.cloud.auth.flows.LeapAuthFlow.from_config_model')
+    def test_revoke(self, flow_factory):
+        flow = flow_factory.return_value
+        token = dict(access_token='123', refresh_token='456')
+        type(flow).token = mock.PropertyMock(return_value=token)
+
+        with self.subTest('dwave auth revoke (default: access-token)'):
+            runner = CliRunner(mix_stderr=False)
+            with runner.isolated_filesystem():
+                result = runner.invoke(cli, ['auth', 'revoke'])
+
+            self.assertEqual(result.exit_code, 0)
+            flow.revoke_token.assert_called_with(
+                token=token['access_token'], token_type_hint='access_token')
+
+        with self.subTest('dwave auth revoke access-token'):
+            runner = CliRunner(mix_stderr=False)
+            with runner.isolated_filesystem():
+                result = runner.invoke(cli, ['auth', 'revoke', 'access-token'])
+
+            self.assertEqual(result.exit_code, 0)
+            flow.revoke_token.assert_called_with(
+                token=token['access_token'], token_type_hint='access_token')
+
+        with self.subTest('dwave auth revoke refresh-token'):
+            runner = CliRunner(mix_stderr=False)
+            with runner.isolated_filesystem():
+                result = runner.invoke(cli, ['auth', 'revoke', 'refresh-token'])
+
+            self.assertEqual(result.exit_code, 0)
+            flow.revoke_token.assert_called_with(
+                token=token['refresh_token'], token_type_hint='refresh_token')
+
+    @mock.patch('dwave.cloud.auth.flows.LeapAuthFlow.from_config_model')
+    def test_revoke_failure_modes(self, flow_factory):
+        flow = flow_factory.return_value
+        type(flow).token = mock.PropertyMock(return_value={})
+
+        with self.subTest('dwave auth revoke: token missing'):
+            runner = CliRunner(mix_stderr=False)
+            with runner.isolated_filesystem():
+                result = runner.invoke(cli, ['auth', 'revoke'])
+
+            self.assertEqual(result.exit_code, 100)
+
+        flow.reset_mock()
+        flow.revoke_token.return_value = False
+        type(flow).token = mock.PropertyMock(return_value=dict(access_token='123'))
+
+        with self.subTest('dwave auth revoke: server-side failure'):
+            runner = CliRunner(mix_stderr=False)
+            with runner.isolated_filesystem():
+                result = runner.invoke(cli, ['auth', 'revoke', 'access-token'])
+
+            self.assertEqual(result.exit_code, 102)
+
     @mock.patch('dwave.cloud.api.resources.LeapAccount.from_config')
     @mock.patch('dwave.cloud.auth.flows.LeapAuthFlow.from_config_model')
     def test_leap_project_ls(self, flow_factory, account_factory):