diff --git a/descope/management/common.py b/descope/management/common.py index 5c1e44b5b..a7f35a791 100644 --- a/descope/management/common.py +++ b/descope/management/common.py @@ -80,6 +80,7 @@ class MgmtV1: # jwt update_jwt_path = "/v1/mgmt/jwt/update" impersonate_path = "/v1/mgmt/impersonate" + stop_impersonation_path = "/v1/mgmt/stop/impersonation" mgmt_sign_in_path = "/v1/mgmt/auth/signin" mgmt_sign_up_path = "/v1/mgmt/auth/signup" mgmt_sign_up_or_in_path = "/v1/mgmt/auth/signup-in" diff --git a/descope/management/jwt.py b/descope/management/jwt.py index 0d01e9b90..a69152a75 100644 --- a/descope/management/jwt.py +++ b/descope/management/jwt.py @@ -59,6 +59,7 @@ def impersonate( validate_consent (bool): Indicate whether to allow impersonation in any case or only if a consent to this operation was granted. customClaims dict: Custom claims to add to JWT tenant_id (str): tenant id to set on DCT claim. + refresh_duration (int): duration in seconds for which the new JWT will be valid Return value (str): A JWT of the impersonated user @@ -86,6 +87,43 @@ def impersonate( pswd=self._auth.management_key, ) return response.json().get("jwt", "") + + def stop_impersonation( + self, + jwt: str, + custom_claims: Optional[dict] = None, + tenant_id: Optional[str] = None, + refresh_duration: Optional[int] = None, + ) -> str: + """ + Stop impersonation and return to the original user + Args: + jwt (str): The impersonation jwt to stop. + customClaims dict: Custom claims to add to JWT + tenant_id (str): tenant id to set on DCT claim. + refresh_duration (int): duration in seconds for which the new JWT will be valid + + Return value (str): A JWT of the actor + + Raise: + AuthException: raised if update failed + """ + if not jwt or jwt == "": + raise AuthException( + 400, ERROR_TYPE_INVALID_ARGUMENT, "jwt cannot be empty" + ) + + response = self._auth.do_post( + MgmtV1.impersonate_path, + { + "jwt": jwt, + "customClaims": custom_claims, + "selectedTenant": tenant_id, + "refreshDuration": refresh_duration, + }, + pswd=self._auth.management_key, + ) + return response.json().get("jwt", "") def sign_in( self, login_id: str, login_options: Optional[MgmtLoginOptions] = None diff --git a/tests/management/test_jwt.py b/tests/management/test_jwt.py index 3baaf1d8f..792a4c4fd 100644 --- a/tests/management/test_jwt.py +++ b/tests/management/test_jwt.py @@ -144,6 +144,50 @@ def test_impersonate(self): timeout=DEFAULT_TIMEOUT_SECONDS, ) + def test_stop_impersonation(self): + client = DescopeClient( + self.dummy_project_id, + self.public_key_dict, + False, + self.dummy_management_key, + ) + + # Test failed flows + with patch("requests.post") as mock_post: + mock_post.return_value.ok = False + self.assertRaises( + AuthException, client.mgmt.jwt.stop_impersonation, "", + ) + + # Test success flow + with patch("requests.post") as mock_post: + network_resp = mock.Mock() + network_resp.ok = True + network_resp.json.return_value = json.loads("""{"jwt": "response"}""") + mock_post.return_value = network_resp + resp = client.mgmt.jwt.stop_impersonation("jwtstr") + self.assertEqual(resp, "response") + expected_uri = f"{common.DEFAULT_BASE_URL}{MgmtV1.impersonate_path}" + mock_post.assert_called_with( + expected_uri, + headers={ + **common.default_headers, + "Authorization": f"Bearer {self.dummy_project_id}:{self.dummy_management_key}", + "x-descope-project-id": self.dummy_project_id, + }, + json={ + "jwt": "jwtstr", + "customClaims": None, + "selectedTenant": None, + "refreshDuration": None, + }, + allow_redirects=False, + verify=True, + params=None, + timeout=DEFAULT_TIMEOUT_SECONDS, + ) + + def test_sign_in(self): client = DescopeClient( self.dummy_project_id,