feat: env var configuration, lint fixes, backoff-retry config for aws…

…-sdk (#51) * feat(aws-sdk): #46 add retry with backoff for aws clients remove unused createSSMClient function * feat(main): allow configuration via env vars prefixed with ECSGO_X * chore: remove unused variables, fix lint errors * tests: remove EC2Client from app_test.go
tedsmitt · Oct 11, 2024 · ff233b1 · ff233b1
1 parent 7260d05
commit ff233b1
Show file tree

Hide file tree

Showing 6 changed files with 55 additions and 77 deletions.
diff --git a/README.md b/README.md
@@ -1,63 +1,70 @@
 # ecsgo
 
-Heavily inspired by the incredibly useful [gossm](https://github.com/gjbae1212/gossm), this tool makes use of the new [ECS ExecuteCommand API](https://aws.amazon.com/blogs/containers/new-using-amazon-ecs-exec-access-your-containers-fargate-ec2/) to connect to running ECS tasks. It provides an interactive prompt to select your cluster, task and container (if only one container in the task it will default to this), and opens a connection to it. You can also use it to port-forward to containers within your tasks.
+Inspired by the incredibly useful [gossm](https://github.com/gjbae1212/gossm), this tool makes use of the [ECS ExecuteCommand API](https://aws.amazon.com/blogs/containers/new-using-amazon-ecs-exec-access-your-containers-fargate-ec2/) to connect to running ECS tasks.
+
+It provides an interactive prompt to select your cluster, task and container (if only one container in the task it will default to this), and opens a connection to it. You can also use it to port-forward to containers within your tasks.
 
 That's it! Nothing fancy.
 
-### Installation
+## Installation
 
-#### MacOS/Homebrew
+### MacOS/Homebrew
 
-```
+```bash
 brew tap tedsmitt/ecsgo
 brew install ecsgo
 ```
 
-#### Linux
+### Linux
 
-```
+```bash
 wget https://github.com/tedsmitt/ecsgo/releases/latest/download/ecsgo_Linux_x86_64.tar.gz
 tar xzf ecsgo_*.tar.gz
 ```
 
 Move the `ecsgo` binary into your `$PATH`
 
-### Pre-requisites
+## Pre-requisites
 
-#### session-manager-plugin
+### session-manager-plugin
 
-This tool makes use of the [session-manager-plugin](https://github.com/aws/session-manager-plugin). For instructions on how to install, please check out https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html.
+This tool makes use of the [session-manager-plugin](https://github.com/aws/session-manager-plugin). For instructions on how to install, please check out [this user guide](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html).
 
-MacOS users can alternatively install this via Homebrew:
+MacOS users can install this via Homebrew if desired
 `brew install --cask session-manager-plugin`
 
-#### Infrastructure
+### Infrastructure
 
 Use [ecs-exec-checker](https://github.com/aws-containers/amazon-ecs-exec-checker) to check for the pre-requisites to use ECS exec.
 
-### Usage
+## Usage
+
+### CLI Args
 
 By default, the tool will prompt you to interactively select which cluster, service, task and container to connect to. You can change the behaviour using the flags detailed below:
-| Long | Short | Description | Default Value |
-| -------------- | ----- | --------------------------------------------------------------------------------------------------------- | -------------------------- |
-| `--cluster` | `-n` | Specify the ECS cluster name | N/A |
-| `--service` | `-s` | Specify the ECS service name | N/A |
-| `--task` | `-t` | Specify the ECS Task ID | N/A |
-| `--container` | `-u` | Specify the container name in the ECS Task (if task only has one container this will selected by default) | N/A |
-| `--cmd` | `-c` | Specify the command to be run on the container (default will change depending on OS family). | `/bin/sh`,`powershell.exe` |
-| `--forward` | `-f` | Port-forward to the container (Remote port will be taken from task/container definitions) | `false` |
-| `--local-port` | `-l` | Specify local port to forward (will prompt if not specified) | N/A |
-| `--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 |
-| `--enable-env` | `-v` | Enable ENV population of cli args | `false` |
 
+| Long                 | Short | Description                                                                                               | Default Value              |
+| -------------------- | ----- | --------------------------------------------------------------------------------------------------------- | -------------------------- |
+| `--cluster`          | `-n`  | Specify the ECS cluster name                                                                              | N/A                        |
+| `--service`          | `-s`  | Specify the ECS service name                                                                              | N/A                        |
+| `--task`             | `-t`  | Specify the ECS Task ID                                                                                   | N/A                        |
+| `--container`        | `-u`  | Specify the container name in the ECS Task (if task only has one container this will selected by default) | N/A                        |
+| `--cmd`              | `-c`  | Specify the command to be run on the container (default will change depending on OS family).              | `/bin/sh`,`powershell.exe` |
+| `--forward`          | `-f`  | Port-forward to the container (Remote port will be taken from task/container definitions)                 | `false`                    |
+| `--local-port`       | `-l`  | Specify local port to forward (will prompt if not specified)                                              | N/A                        |
+| `--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                        |
 
+### Environment variables
 
-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).
+The above options can also be configured via environment variables. Simply export environment variables in the form `ECSGO_<OPT_NAME>`. For example, if you want to set the `--cluster` value, it would be `ECSGO_CLUSTER`, or for the `--aws-endpoint-url` option it would be `ECSGO_AWS_ENDPOINT_URL`.
 
-##### See it in action below
+The tool supports Standard AWS Environment Variables for AWS Client 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).
 
-![ecsgo0 2 0](https://user-images.githubusercontent.com/25430401/114218136-ef8f7b00-9960-11eb-9c3f-b353ae0ff7ca.gif)
+## Example
+
+See it in action below
 
+![ecsgo0 2 0](https://user-images.githubusercontent.com/25430401/114218136-ef8f7b00-9960-11eb-9c3f-b353ae0ff7ca.gif)
diff --git a/cmd/main.go b/cmd/main.go
@@ -69,10 +69,9 @@ Requires pre-existing installation of the session-manager-plugin
 			viper.Set("service", "")
 		}
 
-		if viper.GetBool("enable-env") {
-			viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
-			viper.AutomaticEnv()
-		}
+		viper.SetEnvPrefix("ECSGO")
+		viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
+		viper.AutomaticEnv()
 
 		return nil
 	},
@@ -115,5 +114,4 @@ func init() {
 	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"))
-	viper.BindPFlag("enable-env", rootCmd.PersistentFlags().Lookup("enable-env"))	
 }
diff --git a/internal/app.go b/internal/app.go
@@ -183,7 +183,7 @@ func (e *App) getCluster() {
 		return
 
 	} else {
-		err := errors.New("No clusters found in account or region")
+		err := errors.New("no clusters found in account or region")
 		e.err <- err
 		return
 	}
@@ -383,11 +383,11 @@ func (e *App) getTask() {
 
 	} else {
 		if e.service == "" {
-			fmt.Printf(Red(fmt.Sprintf("There are no running tasks in the cluster %s\n", e.cluster)))
+			fmt.Println(Red(fmt.Sprintf("There are no running tasks in the cluster %s\n", e.cluster)))
 			e.input <- "getCluster"
 			return
 		} else {
-			fmt.Printf(Red(fmt.Sprintf("\nThere are no running tasks for the service %s in cluster %s\n", e.service, e.cluster)))
+			fmt.Println(Red(fmt.Sprintf("\nThere are no running tasks for the service %s in cluster %s\n", e.service, e.cluster)))
 			e.input <- "getService"
 			return
 		}
@@ -406,7 +406,7 @@ func (e *App) getContainer() {
 				return
 			}
 		}
-		fmt.Printf(Red(fmt.Sprintf("\nSupplied container with name %s not found in task %s, cluster %s\n", cliArg, *e.task.TaskArn, e.cluster)))
+		fmt.Println(Red(fmt.Sprintf("\nSupplied container with name %s not found in task %s, cluster %s\n", cliArg, *e.task.TaskArn, e.cluster)))
 	}
 
 	if len(e.task.Containers) > 1 {
@@ -437,7 +437,7 @@ func (e *App) getContainer() {
 func (e *App) getContainerOS() {
 	// Get associated task definition and determine OS family if EC2 launch-type
 	if e.task.LaunchType == "EC2" {
-		family, err := getPlatformFamily(e.client, e.cluster, e.task)
+		family, err := getPlatformFamily(e.client, e.task)
 		if err != nil {
 			e.err <- err
 			return

diff --git a/internal/app_test.go b/internal/app_test.go
@@ -7,7 +7,6 @@ import (
 	"testing"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
-	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/aws-sdk-go-v2/service/ecs"
 	ecsTypes "github.com/aws/aws-sdk-go-v2/service/ecs/types"
 	"github.com/stretchr/testify/assert"
@@ -54,18 +53,6 @@ func (m ECSClientMock) ExecuteCommand(ctx context.Context, params *ecs.ExecuteCo
 	return m.ExecuteCommandMock(ctx, params, optFns...)
 }
 
-type MockEC2API struct {
-	ec2.Client            // embedding of the interface is needed to skip implementation of all methods
-	DescribeInstancesMock func(input *ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error)
-}
-
-func (m *MockEC2API) DescribeInstances(input *ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error) {
-	if m.DescribeInstancesMock != nil {
-		return m.DescribeInstancesMock(input)
-	}
-	return nil, nil
-}
-
 // CreateMockApp initialises a new App struct and takes a MockClient as an argument - only used in tests
 func CreateMockApp(c ECSClient) *App {
 	e := &App{

diff --git a/internal/internal.go b/internal/internal.go
@@ -7,20 +7,20 @@ import (
 	"os/exec"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/retry"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	"github.com/aws/aws-sdk-go-v2/service/ecs"
 	ecsTypes "github.com/aws/aws-sdk-go-v2/service/ecs/types"
-	"github.com/aws/aws-sdk-go-v2/service/ssm"
 	"github.com/fatih/color"
 	"github.com/spf13/viper"
 )
 
 var (
-	region   string
-	endpoint string
+	region string
 
 	Red     = color.New(color.FgRed).SprintFunc()
 	Magenta = color.New(color.FgMagenta).SprintFunc()
@@ -49,6 +49,9 @@ func createEcsClient() *ecs.Client {
 	cfg, err := config.LoadDefaultConfig(context.Background(),
 		config.WithSharedConfigProfile(viper.GetString("profile")),
 		config.WithRegion(region),
+		config.WithRetryer(func() aws.Retryer {
+			return retry.AddWithMaxBackoffDelay(retry.NewStandard(), time.Second*1)
+		}),
 	)
 	if err != nil {
 		panic(err)
@@ -69,6 +72,9 @@ func createEc2Client() *ec2.Client {
 	cfg, err := config.LoadDefaultConfig(context.Background(),
 		config.WithSharedConfigProfile(viper.GetString("profile")),
 		config.WithRegion(region),
+		config.WithRetryer(func() aws.Retryer {
+			return retry.AddWithMaxBackoffDelay(retry.NewStandard(), time.Second*1)
+		}),
 	)
 	if err != nil {
 		panic(err)
@@ -78,29 +84,9 @@ func createEc2Client() *ec2.Client {
 	return client
 }
 
-func createSSMClient() *ssm.Client {
-	region := viper.GetString("region")
-	getCustomAWSEndpoint := func(o *ssm.Options) {
-		endpointUrl := viper.GetString("aws-endpoint-url")
-		if endpointUrl != "" {
-			o.BaseEndpoint = aws.String(endpointUrl)
-		}
-	}
-	cfg, err := config.LoadDefaultConfig(context.Background(),
-		config.WithSharedConfigProfile(viper.GetString("profile")),
-		config.WithRegion(region),
-	)
-	if err != nil {
-		panic(err)
-	}
-	client := ssm.NewFromConfig(cfg, getCustomAWSEndpoint)
-
-	return client
-}
-
 // getPlatformFamily checks an ECS tasks properties to see if the OS can be derived from its properties, otherwise
 // it will check the container instance itself to determine the OS.
-func getPlatformFamily(client ECSClient, clusterName string, task *ecsTypes.Task) (string, error) {
+func getPlatformFamily(client ECSClient, task *ecsTypes.Task) (string, error) {
 	taskDefinition, err := client.DescribeTaskDefinition(context.TODO(), &ecs.DescribeTaskDefinitionInput{
 		TaskDefinition: task.TaskDefinitionArn,
 	})

diff --git a/internal/internal_test.go b/internal/internal_test.go
@@ -75,7 +75,7 @@ func TestGetPlatformFamily(t *testing.T) {
 
 	for _, c := range cases {
 		client := c.client(t)
-		res, _ := getPlatformFamily(client, c.cluster, c.task)
+		res, _ := getPlatformFamily(client, c.task)
 		if ok := assert.Equal(t, c.expected, res); ok != true {
 			fmt.Printf("%s FAILED\n", c.name)
 		}