support GovCloud IAM role ARN

Author: Michael Kania

.circleci/config.yml

@@ -1,7 +1,7 @@
 version: 2.1
 
 references:
-  circleci-docker-primary: &circleci-docker-primary trussworks/circleci-docker-primary:d7f27acdab24f86297220a43f704f23f2bab667d
+  circleci-docker-primary: &circleci-docker-primary trussworks/circleci-docker-primary:c542b22c7fb95db0a1bbe043928a457ae6fbeaca
 
 jobs:
   test:

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: git://github.com/golangci/golangci-lint
-    rev: v1.24.0
+    rev: v1.26.0
     hooks:
       - id: golangci-lint
 
@@ -17,6 +17,6 @@ repos:
       - id: trailing-whitespace
 
   - repo: git://github.com/igorshubovych/markdownlint-cli
-    rev: v0.22.0
+    rev: v0.23.0
     hooks:
       - id: markdownlint

cmd/main.go

@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"io/ioutil"
-
 	"log"
 	"os"
 
@@ -12,6 +11,7 @@ import (
 	"github.com/99designs/keyring"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/endpoints"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/iam"
 	"github.com/aws/aws-sdk-go/service/sts"
@@ -453,6 +453,13 @@ func checkExistingAWSProfile(profileName string, config *vault.Config) error {
 	return nil
 }
 
+func getPartition(region string) string {
+	if partition, ok := endpoints.PartitionForRegion(endpoints.DefaultPartitions(), region); ok {
+		return partition.ID()
+	}
+	return "aws"
+}
+
 func main() {
 	// parse command line flags
 	var options cliOptions
@@ -468,8 +475,10 @@ func main() {
 	// initialize things
 	profile := vault.Profile{
 		Name: options.AwsProfile,
-		RoleARN: fmt.Sprintf("arn:aws:iam::%v:role/%v",
-			options.AwsAccountID, options.Role),
+		RoleARN: fmt.Sprintf("arn:%s:iam::%d:role/%s",
+			getPartition(options.AwsRegion),
+			options.AwsAccountID,
+			options.Role),
 		Region: options.AwsRegion,
 	}
 

cmd/main_test.go

@@ -61,3 +61,12 @@ func TestGenerateQrCode(t *testing.T) {
 	err = generateQrCode("otpauth://totp/super@top?secret=secret", tempFile)
 	assert.NoError(t, err)
 }
+
+func TestGetPartition(t *testing.T) {
+	commPartition := getPartition("us-west-2")
+	assert.Equal(t, commPartition, "aws")
+	govPartition := getPartition("us-gov-west-1")
+	assert.Equal(t, govPartition, "aws-us-gov")
+	unknownPartition := getPartition("aws-under-the-sea")
+	assert.Equal(t, unknownPartition, "aws")
+}