From 68e578de11b61d77432706af76387b80538d7f1b Mon Sep 17 00:00:00 2001
From: Mike Forkin <mforkin@leolabs.space>
Date: Thu, 12 Sep 2024 14:09:09 -0700
Subject: [PATCH] feat: Allow endpoint url to be passed through to aws clients

---
 README.md            | 1 +
 cmd/main.go          | 6 ++++++
 internal/internal.go | 9 ++++++---
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index de6ddd2..57b8432 100644
--- a/README.md
+++ b/README.md
@@ -50,6 +50,7 @@ By default, the tool will prompt you to interactively select which cluster, serv
 | `--profile` | `-p` | Specify the profile to load the credentials | `default` |
 | `--region` | `-r` | Specify the AWS region to run in | N/A |
 | `--quiet` | `-q` | Disable output detailing the Cluster/Service/Task information | `false` |
+| `--aws-endpoint-url` | `-e` | Specify the AWS endpoint used for all service requests | N/A |
 
 The tool also supports AWS Config/Environment Variables for configuration. If you aren't familiar with working on AWS via the CLI, you can read more about how to configure your environment [here](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html).
 
diff --git a/cmd/main.go b/cmd/main.go
index 3608d4f..dc5726e 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -18,6 +18,7 @@ package main
 import (
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -94,6 +95,10 @@ func init() {
 	rootCmd.PersistentFlags().BoolP("forward", "f", false, "Port Forward")
 	rootCmd.PersistentFlags().StringP("local-port", "l", "", "Local port for use with port forwarding")
 	rootCmd.PersistentFlags().BoolP("quiet", "q", false, "Do not print cluster and container information")
+	rootCmd.PersistentFlags().StringP("aws-endpoint-url", "e", "", "AWS Endpoint Url")
+
+	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
+	viper.AutomaticEnv()
 
 	viper.BindPFlag("cmd", rootCmd.PersistentFlags().Lookup("cmd"))
 	viper.BindPFlag("profile", rootCmd.PersistentFlags().Lookup("profile"))
@@ -105,4 +110,5 @@ func init() {
 	viper.BindPFlag("forward", rootCmd.PersistentFlags().Lookup("forward"))
 	viper.BindPFlag("local-port", rootCmd.PersistentFlags().Lookup("local-port"))
 	viper.BindPFlag("quiet", rootCmd.PersistentFlags().Lookup("quiet"))
+	viper.BindPFlag("aws-endpoint-url", rootCmd.PersistentFlags().Lookup("aws-endpoint-url"))
 }
diff --git a/internal/internal.go b/internal/internal.go
index 66b38a7..8c1b6a9 100644
--- a/internal/internal.go
+++ b/internal/internal.go
@@ -40,8 +40,9 @@ func createOpts(opts []string) []string {
 
 func createEcsClient() *ecs.ECS {
 	region := viper.GetString("region")
+	endpointUrl := viper.GetString("aws-endpoint-url")
 	sess := session.Must(session.NewSessionWithOptions(session.Options{
-		Config:            aws.Config{Region: aws.String(region)},
+		Config:            aws.Config{Region: aws.String(region), Endpoint: aws.String(endpointUrl)},
 		Profile:           viper.GetString("profile"),
 		SharedConfigState: session.SharedConfigEnable,
 	}))
@@ -52,8 +53,9 @@ func createEcsClient() *ecs.ECS {
 
 func createEc2Client() *ec2.EC2 {
 	region := viper.GetString("region")
+	endpointUrl := viper.GetString("aws-endpoint-url")
 	sess := session.Must(session.NewSessionWithOptions(session.Options{
-		Config:            aws.Config{Region: aws.String(region)},
+		Config:            aws.Config{Region: aws.String(region), Endpoint: aws.String(endpointUrl)},
 		Profile:           viper.GetString("profile"),
 		SharedConfigState: session.SharedConfigEnable,
 	}))
@@ -64,8 +66,9 @@ func createEc2Client() *ec2.EC2 {
 
 func createSSMClient() *ssm.SSM {
 	region := viper.GetString("region")
+	endpointUrl := viper.GetString("aws-endpoint-url")
 	sess := session.Must(session.NewSessionWithOptions(session.Options{
-		Config:            aws.Config{Region: aws.String(region)},
+		Config:            aws.Config{Region: aws.String(region), Endpoint: aws.String(endpointUrl)},
 		Profile:           viper.GetString("profile"),
 		SharedConfigState: session.SharedConfigEnable,
 	}))