Initial commit

Signed-off-by: Jordan Jacobelli <[email protected]>

.github/workflows/pr.yml

@@ -0,0 +1,22 @@
+name: pr
+
+on:
+  pull_request:
+
+jobs:
+  run:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Setup Go
+        uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+      - name: Check gofmt
+        run: make gofmt-verify
+      - name: Run golangci-lint
+        uses: golangci/golangci-lint-action@v5
+        with:
+          version: latest
+          args: --timeout=5m

.github/workflows/release.yml

@@ -0,0 +1,14 @@
+name: release
+on:
+  push:
+    tags:
+      - "v*"
+permissions:
+  contents: write
+
+jobs:
+  release:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: cli/gh-extension-precompile@v1

.gitignore

@@ -0,0 +1,2 @@
+/gh-nv-gha-aws
+/gh-nv-gha-aws.exe

Makefile

@@ -0,0 +1,25 @@
+BIN_NAME = gh-nv-gha-aws
+VERSION = $(shell git describe --tags --dirty --always)
+BUILD_FLAGS = -tags osusergo,netgo \
+              -ldflags "-s -extldflags=-static -X main.version=$(VERSION)"
+
+build:
+	go build -o $(BIN_NAME) $(BUILD_FLAGS)
+
+check: gofmt-verify ci-lint
+
+gofmt:
+	@gofmt -w -l $$(find . -name '*.go')
+
+gofmt-verify:
+	@out=`gofmt -w -l -d $$(find . -name '*.go')`; \
+		if [ -n "$$out" ]; then \
+		echo "$$out"; \
+		exit 1; \
+		fi
+
+ci-lint:
+	@docker run --pull always --rm -v $(PWD):/app -w /app golangci/golangci-lint:latest golangci-lint run
+
+clean:
+	@rm -f $(BIN_NAME)

README.md

@@ -0,0 +1,14 @@
+# gh-nv-gha-aws
+
+`gh-nv-gha-aws` is a `gh` extension that allows users to obtain temporary AWS credentials for preconfigured IAM roles based on GitHub organization or team membership.
+
+## Steps to install
+
+1. Please ensure that you have the `gh` CLI tool [installed](https://docs.github.com/en/github-cli/github-cli/quickstart).
+
+2. Login and authenticate with `gh auth login`. This is required to ensure you have correct credentials to receive the AWS Credentials.
+
+3. Run `gh extension install nv-gha-aws`
+
+More information about all of the available flags and their associated usage is available when running `gh nv-gha-aws --help`
+

cmd/org.go

@@ -0,0 +1,122 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"os/signal"
+	"syscall"
+
+	"github.com/spf13/cobra"
+
+	"github.com/nv-gha-runners/gh-nv-gha-aws/pkg/aws"
+	"github.com/nv-gha-runners/gh-nv-gha-aws/pkg/gh"
+	"github.com/nv-gha-runners/gh-nv-gha-aws/pkg/jwt"
+)
+
+var orgCmd = &cobra.Command{
+	Use:   "org",
+	Short: "Receive AWS Credentials by providing an organization name",
+	Args:  cobra.ExactArgs(1),
+	RunE: func(command *cobra.Command, args []string) error {
+		ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+		defer cancel()
+
+		orgName := args[0]
+
+		idpUrl, err := command.Flags().GetString("idp-url")
+		if err != nil {
+			return fmt.Errorf("failed to get --idp-url flag: %w", err)
+		}
+
+		aud, err := command.Flags().GetString("aud")
+		if err != nil {
+			return fmt.Errorf("failed to get --aud flag: %w", err)
+		}
+
+		roleArn, err := command.Flags().GetString("role-arn")
+		if err != nil {
+			return fmt.Errorf("failed to get --role-arn flag: %w", err)
+		}
+
+		duration, err := command.Flags().GetInt32("duration")
+		if err != nil {
+			return fmt.Errorf("failed to get --duration flag: %w", err)
+		}
+
+		profile, err := command.Flags().GetString("profile")
+		if err != nil {
+			return fmt.Errorf("failed to get --profile flag: %w", err)
+		}
+
+		output, err := command.Flags().GetString("output")
+		if err != nil {
+			return fmt.Errorf("failed to get --output flag: %w", err)
+		}
+
+		write, err := command.Flags().GetBool("write")
+		if err != nil {
+			return fmt.Errorf("failed to get --write flag: %w", err)
+		}
+
+		file, err := command.Flags().GetString("file")
+		if err != nil {
+			return fmt.Errorf("failed to get --file flag: %w", err)
+		}
+
+		ghToken, err := gh.GetGHToken()
+		if err != nil {
+			return err
+		}
+
+		ghClient, err := gh.NewClient(ghToken)
+		if err != nil {
+			return fmt.Errorf("failed to create GH client: %w", err)
+		}
+
+		username, err := ghClient.GetUsername()
+		if err != nil {
+			return fmt.Errorf("failed to get username: %w", err)
+		}
+
+		orgID, err := ghClient.GetOrgID(orgName)
+		if err != nil {
+			return fmt.Errorf("failed to get org ID: %w", err)
+		}
+
+		jwt, err := jwt.GetOrgJWT(&jwt.JWTInputs{
+			Audience: aud,
+			GHToken:  ghToken,
+			IDPUrl:   idpUrl,
+		}, orgID)
+		if err != nil {
+			return fmt.Errorf("failed to get org JWT: %w", err)
+		}
+
+		creds, err := aws.GetCreds(ctx, &aws.GetCredsInput{
+			Duration: duration,
+			JWT:      jwt,
+			Profile:  profile,
+			RoleArn:  roleArn,
+			Username: username,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to fet AWS credentials: %w", err)
+		}
+
+		if err = creds.Print(output); err != nil {
+			return fmt.Errorf("failed to print credentials: %w", err)
+		}
+
+		if write {
+			if err = creds.Write(file); err != nil {
+				return fmt.Errorf("failed to write credentials file: %w", err)
+			}
+		}
+
+		return nil
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(orgCmd)
+}

cmd/root.go

@@ -0,0 +1,67 @@
+package cmd
+
+import (
+	"errors"
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+)
+
+const (
+	programName = "nv-gha-aws"
+)
+
+var rootCmd = &cobra.Command{
+	Use:          programName,
+	Short:        "A GitHub CLI Extension Tool to receive AWS Credentials",
+	SilenceUsage: true,
+}
+
+func Execute(version string) {
+	rootCmd.Version = version
+	err := rootCmd.Execute()
+	if err != nil {
+		os.Exit(1)
+	}
+}
+
+func init() {
+	rootCmd.PersistentFlags().String("role-arn", "", "Role ARN ")
+	_ = rootCmd.MarkFlagRequired("role-arn")
+
+	rootCmd.PersistentFlags().String("idp-url", "https://token.gha-runners.nvidia.com", "Identity Provider URL")
+	rootCmd.PersistentFlags().String("aud", "sts.amazonaws.com", "Audience of Web Identity Token")
+	rootCmd.PersistentFlags().Int32P("duration", "d", 43200, "The maximum session duration with the temporary AWS Credentials in seconds")
+	rootCmd.PersistentFlags().StringP("output", "o", "shell", "Output format of credentials in one of: shell, json, or creds-file format")
+	rootCmd.PersistentFlags().BoolP("write", "w", false, "Specifies if Credentials should be written to AWS Credentials file")
+	rootCmd.PersistentFlags().StringP("file", "f", "$HOME/.aws/credentials", "File path to write AWS Credentials")
+	rootCmd.PersistentFlags().StringP("profile", "p", "default", "Profile where credentials should be written to in AWS Credentials File")
+
+	rootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		writeFlag, err := cmd.Flags().GetBool("write")
+		if err != nil {
+			return fmt.Errorf("failed to get --write flag: %w", err)
+		}
+
+		printFormatFlag, err := cmd.Flags().GetString("output")
+		if err != nil {
+			return fmt.Errorf("failed to get --output flag: %w", err)
+		}
+
+		// ensures that the writeFlag must be set in order to set the file flag
+		if cmd.Flags().Changed("file") && !writeFlag {
+			return errors.New("the write flag must be set if specifying a file path")
+		}
+
+		if printFormatFlag != "creds-file" && cmd.Flags().Changed("profile") {
+			return errors.New("the profile can only be set if the output flag is set to creds-file")
+		}
+
+		if printFormatFlag != "creds-file" && writeFlag {
+			return errors.New("the write flag can only be set if the output flag is set to creds-file")
+		}
+
+		return nil
+	}
+}