Okta for AWS - local setup and Go wrapper around okta-aws-cli.jar executable.

.bash_profile

. "${HOME}/.java_version"
. "${HOME}/.okta_functions"

.java_version

JAVA_VERSION='1.8.0_191'
JAVA_HOME="${HOME}/.binaries/java/${JAVA_VERSION}"
JDK_HOME="$JAVA_HOME"
JRE_HOME="${JAVA_HOME}/jre"
JAVA_TOOL_OPTIONS=

export PATH="${JAVA_HOME}/bin:${PATH}"
export JAVA_HOME JDK_HOME JRE_HOME JAVA_TOOL_OPTIONS

.okta_functions

okta-aws() {
    with-okta "aws --profile $1" "$@"
}

okta-sls() {
    with-okta "sls --stage $1" "$@"
}

okta-aws-<ROLE NAME>() {
    export OKTA_AWS_ROLE_TO_ASSUME="arn:aws:iam::<ACCOUNt ID>:role/<ROLE NAME>"
    okta-login <ROLE NAME>
}

main.go

// with-okta-env
// $ go get -v github.com/go-ini/ini
// $ go get -v github.com/mitchellh/go-homedir
// $ GOOS=linux go build -o with-okta-env .
package main

import (
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/go-ini/ini"
	homedir "github.com/mitchellh/go-homedir"
)

const (
	OktaEnvJarFile     = "OKTA_JAR_FILE"
	OktaDefaultJarFile = ".okta/okta-aws-cli.jar"
	OktaJavaClass      = "com.okta.tools.WithOkta"
	AWSCredsFile       = ".aws/credentials"
)

var AWSConfigMap = map[string][]string{
	"aws_access_key_id": {
		"AWS_ACCESS_KEY_ID",
	},
	"aws_secret_access_key": {
		"AWS_SECRET_ACCESS_KEY",
	},
	"aws_session_token": {
		"AWS_SESSION_TOKEN",
		"AWS_SECURITY_TOKEN",
	},
	"profile": {
		"AWS_DEFAULT_PROFILE",
		"AWS_PROFILE",
	},
	"region": {
		"AWS_DEFAULT_REGION",
		"AWS_REGION",
	},
}

type environ []string

func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], fmt.Sprintf("%s=", key)) {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

func (e *environ) Set(key, value string) {
	e.Unset(key)
	*e = append(*e, fmt.Sprintf("%s=%s", key, value))
}

func checkFileExist(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func checkCmdExist(name string) bool {
	cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("command -v %s", name))
	err := cmd.Run()
	return err == nil
}

func abort(msg string, args ...interface{}) {
	_, err := fmt.Fprintf(os.Stderr, msg, args...)
	if err != nil {
		panic(err)
	}
	os.Exit(1)
}

func main() {
	userHomeDir, err := homedir.Dir()
	if err != nil {
		abort("Unable to determine home directory for current user: %s\n", err)
	}

	if !checkCmdExist("java") {
		abort("Unable to locate \"java\" executable in the current environment. Java not installed?\n")
	}

	oktaJarFile := os.Getenv(OktaEnvJarFile)
	if oktaJarFile == "" {
		oktaJarFile = path.Join(userHomeDir, OktaDefaultJarFile)
	}
	if !checkFileExist(oktaJarFile) {
		abort("Unable to locate Okta JAR file at %s\n", oktaJarFile)
	}

	if len(os.Args[1:]) < 2 {
		abort("Not enough arguments. Usage: %s <PROFILE> <COMMAND> [<ARGUMENTS> ...]\n", os.Args[0])
	}
	awsProfile := os.Args[1]

	env := environ(os.Environ())
	env.Set("OKTA_PROFILE", awsProfile)

	cmdArgs := []string{"-classpath", oktaJarFile, OktaJavaClass, "true"}

	cmd := exec.Command("java", cmdArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = env

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigChan
		if cmd.Process != nil {
			errSig := cmd.Process.Signal(sig)
			if errSig != nil {
				abort("An error occurred while forwarding signal: %s\n", errSig)
			}
		}
		signal.Stop(sigChan)
	}()

	err = cmd.Run()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			waitStatus := exitErr.Sys().(syscall.WaitStatus)
			os.Exit(waitStatus.ExitStatus())
		}
		abort("An unknown error occurred while running Java: %s\n", err)
	}

	for _, v := range AWSConfigMap {
		for _, envName := range v {
			env.Unset(envName)
		}
	}
	env.Unset("OKTA_PROFILE")

	credsFile := filepath.Join(userHomeDir, AWSCredsFile)
	if !checkFileExist(credsFile) {
		abort("Unable to locate AWS credentials files at %s\n", credsFile)

	}
	ini, err := ini.Load(credsFile)
	if err != nil {
		abort("Unable to load AWS credentials file: %s\n", err)
	}
	section, err := ini.GetSection(awsProfile)
	if err != nil {
		abort("Unable to load AWS profile %q from credentials file: %s\n", awsProfile, err)
	}

	for k, v := range AWSConfigMap {
		if section.HasKey(k) {
			for _, envName := range v {
				env.Set(envName, section.Key(k).Value())
			}
		}
	}

	cmdName := os.Args[2]
	if !checkCmdExist(cmdName) {
		abort("Unable to locate %q executable in the current environment.\n", cmdName)
	}

	cmd = exec.Command(cmdName, os.Args[3:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = env

	signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigChan
		if cmd.Process != nil {
			errSig := cmd.Process.Signal(sig)
			if errSig != nil {
				abort("An error occurred while forwarding signal: %s\n", errSig)
			}
		}
		signal.Stop(sigChan)
	}()

	err = cmd.Run()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			waitStatus := exitErr.Sys().(syscall.WaitStatus)
			os.Exit(waitStatus.ExitStatus())
		}
		abort("An unknown error occurred while running command %q: %s\n", cmdName, err)
	}
}

okta-aws-cli

#!/bin/bash

set -e

exec "$(dirname $0)/okta-run-class.sh" com.okta.tools.awscli "$@"

okta-list-roles

#!/bin/bash

set -e

exec "$(dirname $0)/okta-run-class.sh" com.okta.tools.ListRoles

okta-login

#!/bin/bash

set -e

profile="$1"
shift

export OKTA_PROFILE="$profile"

"$(dirname $0)/okta-run-class.sh" com.okta.tools.WithOkta true

okta-logout

#!/bin/bash

set -e

"$(dirname $0)/okta-run-class.sh" com.okta.tools.WithOkta logout

okta-reset

#!/bin/bash

set -e

"$(dirname $0)/okta-run-class.sh" com.okta.tools.WithOkta logout &>/dev/null

unset $(env | grep -E '^(AWS_|OKTA_)' | awk -F'=' '{ print $1 }')

rm -f "${HOME}/.okta/cookies.properties" \
      "${HOME}/.okta/profiles"

okta-run-class.sh

#!/bin/bash

set -e

export PATH='/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin'

"${JAVA_HOME}/bin/java" -classpath "${HOME}/.okta/okta-aws-cli.jar" "$@"

with-okta

#!/bin/bash

set -e

command="$1"
shift
profile="$1"
shift

export OKTA_PROFILE="$profile"

exec "$(dirname $0)/okta-run-class.sh" com.okta.tools.WithOkta $command "$@"