ECR Docker image puller / installer

installer.go

package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecr"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/client"
)

func assumeRoleWithRetry(role string) (*sts.AssumeRoleOutput, *session.Session, error) {
	fmt.Println("Creating session with default credentials...")
	sess, err := session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
		Config: aws.Config{
			Region: aws.String("eu-central-1"),
		},
	})

	if err != nil {
		fmt.Println("Failed to create session with default credentials. Please enter AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY.")

		var awsAccessKeyID, awsSecretAccessKey string
		fmt.Print("Enter AWS_ACCESS_KEY_ID: ")
		fmt.Scanln(&awsAccessKeyID)

		fmt.Print("Enter AWS_SECRET_ACCESS_KEY: ")
		fmt.Scanln(&awsSecretAccessKey)

		fmt.Println("Creating session with provided credentials...")

		sess, _ = session.NewSession(&aws.Config{
			Region:      aws.String("eu-central-1"),
			Credentials: credentials.NewStaticCredentials(awsAccessKeyID, awsSecretAccessKey, ""),
		})
	}

	fmt.Println("Assuming role...")
	svc := sts.New(sess)
	input := &sts.AssumeRoleInput{
		RoleArn:         aws.String(role),
		RoleSessionName: aws.String("sessionName"),
	}

	result, err := svc.AssumeRole(input)
	return result, sess, err
}

func main() {
	const image = "XXXXXXXX.dkr.ecr.XXXXXXXX.amazonaws.com/XXXXXXXX"
	const role = "arn:aws:iam::XXXXXXXXXXX:role/XXXXXXXXX"

	result, sess, err := assumeRoleWithRetry(role)
	if err != nil {
		fmt.Println("Failed to assume role:", err)
		os.Exit(-1)
	} else {
		fmt.Println("Assumed role successfully")
	}

	fmt.Println("Getting ECR token...")
	svcECR := ecr.New(sess, &aws.Config{Credentials: credentials.NewStaticCredentials(*result.Credentials.AccessKeyId, *result.Credentials.SecretAccessKey, *result.Credentials.SessionToken)})

	ecrInput := &ecr.GetAuthorizationTokenInput{}
	ecrResult, _ := svcECR.GetAuthorizationToken(ecrInput)

	authData, _ := base64.StdEncoding.DecodeString(*ecrResult.AuthorizationData[0].AuthorizationToken)
	authParts := strings.Split(string(authData), ":")

	fmt.Println("Logging in to Docker / ECR...")
	cli, _ := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
	authConfig := types.AuthConfig{
		Username: authParts[0],
		Password: authParts[1],
	}
	encodedJSON, _ := json.Marshal(authConfig)
	authStr := base64.URLEncoding.EncodeToString(encodedJSON)

	fmt.Println("Pulling", image)
	out, err := cli.ImagePull(context.Background(), image, types.ImagePullOptions{RegistryAuth: authStr})
	if err != nil {
		fmt.Println("Failed to pull image:", err)
		os.Exit(-1)
	}
	defer out.Close()
	io.Copy(os.Stdout, out)
	fmt.Println("Image pulled successfully")
}