#69 - Fix STS endpoint for China

Author: ivan-khvostishkov

sagemaker_ssh_helper/aws.py

@@ -15,6 +15,11 @@ def get_console_domain(self):
             return f"{self.region}.console.amazonaws-us-gov.com"
         return f"{self.region}.console.aws.amazon.com"
 
+    def get_sts_endpoint(self):
+        if self.region.startswith("cn-"):
+            return "https://sts.{}.amazonaws.com.cn".format(self.region)
+        return "https://sts.{}.amazonaws.com".format(self.region)
+
     @staticmethod
     def is_arn(arn: str):
         """

sagemaker_ssh_helper/wrapper.py

@@ -73,7 +73,8 @@ def _augment(self):
     def _augment_env(self, env):
         if self.local_user_id is None:
             region = self.sagemaker_session.boto_region_name
-            endpoint_url = "https://sts.{}.amazonaws.com".format(region)
+            aws = AWS(region)
+            endpoint_url = aws.get_sts_endpoint()
             caller_id = boto3.client("sts", region_name=region, endpoint_url=endpoint_url).get_caller_identity()
             user_id = caller_id.get('UserId')
         else: