From 807fe1c44a45ddd5145e5a7254c39976c4d260db Mon Sep 17 00:00:00 2001
From: Anders Pedersen <37378030+anderspede@users.noreply.github.com>
Date: Tue, 3 Oct 2023 10:48:31 +0200
Subject: [PATCH] Add support for AWS China (#14)

* Set partition based on region

* Add China support for launching the console

* Fix syntax error

* strings should be quoted

* Update assume-aws-sso-role

Co-authored-by: Jonas Holst Damtoft <56115984+jhd-lego@users.noreply.github.com>

---------

Co-authored-by: Jonas Holst Damtoft <56115984+jhd-lego@users.noreply.github.com>
---
 assume-aws-sso-role | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/assume-aws-sso-role b/assume-aws-sso-role
index 6f4e8fd..bef1c99 100755
--- a/assume-aws-sso-role
+++ b/assume-aws-sso-role
@@ -31,6 +31,7 @@ aws_role=
 aws_account=
 aws_role=
 aws_session_duration=3600
+aws_partition="aws"
 
 mkdir -p "${credentials_path}"
 
@@ -94,7 +95,7 @@ role_name_from_role_arn() {
 }
 
 set_role_arn_from_account_and_role() {
-  aws_role_arn="arn:aws:iam::${aws_account}:role/${aws_role}"
+  aws_role_arn="arn:${aws_partition}:iam::${aws_account}:role/${aws_role}"
 }
 
 set_account_and_role_from_role_arn() {
@@ -105,6 +106,7 @@ set_account_and_role_from_role_arn() {
 role_info_from_profile() {
   if [ -z "${aws_region}" ]; then
     aws_region=$(aws configure get region --profile "${aws_profile}")
+    get_partition_from_region
   fi
   if [ -z "${aws_region}" ]; then
     echo_maybe "Region not found in profile ${aws_profile}."
@@ -118,6 +120,15 @@ role_info_from_profile() {
   set_account_and_role_from_role_arn
 }
 
+get_partition_from_region() {
+  if [[ ${aws_region} = cn-* ]]; then
+    aws_partition="aws-cn"
+    federation_destination_url="https://console.amazonaws.cn"
+    aws_federation_signin_url="https://signin.amazonaws.cn/federation"
+    logout_url="https://signin.amazonaws.cn/oauth?Action=logout"
+  fi
+}
+
 url_encode() {
   echo "$1" | jq "@uri" -jR
 }
@@ -307,6 +318,7 @@ if [ -n "${aws_profile}" ]; then
     ${exit_cmd} 3
   fi
 else
+  get_partition_from_region
   set_role_arn_from_account_and_role
 fi
 
@@ -327,7 +339,7 @@ echo_maybe "Region is ${aws_region}"
 echo_maybe "Role arn is ${aws_role_arn}"
 echo_maybe "Aws account is ${aws_account}"
 
-aws_provider="arn:aws:iam::${aws_account}:saml-provider/AzureAD"
+aws_provider="arn:${aws_partition}:iam::${aws_account}:saml-provider/AzureAD"
 
 # Obtain token for OIDCtoSAML endpoint
 echo_maybe "Getting access token"