Stop impersonation call (#565)

related to descope/etc#9294
+ tests

Author: aviadl

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"

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

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,