Initial Commit

README.md

@@ -0,0 +1,23 @@
+# Locksmith in Go
+
+Configure your `~/.aws/credentials` with your AWS and Beagle credentials:
+
+```
+[default]
+aws_access_key_id = AKIAXXXXXXXXXXXXXXXX
+aws_secret_access_key = XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+mfa_serial = XXXXXXXXXXXXXXXXXXXXXXX
+beagle_url = https://beagle.sentiampc.com/api/v1/bookmarks
+beagle_pass = XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+```
+
+Install:
+
+```
+export GOPATH=~/go
+export GOBIN=$GOPATH/bin
+
+go get -u github.com/sentialabs/locksmith-go/locksmith
+```
+
+Now, add `~/go/bin` to your path. And start `locksmith`!

locksmith/main.go

@@ -0,0 +1,225 @@
+package main
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"os/exec"
+	"regexp"
+	"sort"
+	"strings"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/sts"
+	"github.com/captainsafia/go-user-shell"
+	"github.com/manifoldco/promptui"
+	"gopkg.in/ini.v1"
+	"gopkg.in/mattes/go-expand-tilde.v1"
+)
+
+type Bookmarks struct {
+	Links struct {
+		Parent struct {
+			Href string `json:"href"`
+		} `json:"parent"`
+		First struct {
+			Href string `json:"href"`
+		} `json:"first"`
+		Self struct {
+			Href string `json:"href"`
+		} `json:"self"`
+		Last struct {
+			Href string `json:"href"`
+		} `json:"last"`
+	} `json:"_links"`
+	Bookmarks []struct {
+		Links struct {
+			Parent struct {
+				Href string `json:"href"`
+			} `json:"parent"`
+			Self struct {
+				Href string `json:"href"`
+			} `json:"self"`
+		} `json:"_links"`
+		ID            int    `json:"id"`
+		RoleName      string `json:"role_name"`
+		Name          string `json:"name"`
+		AccountNumber string `json:"account_number"`
+		AvatarURL     string `json:"avatar_url"`
+	} `json:"bookmarks"`
+	TotalCount int `json:"total_count"`
+	TotalPages int `json:"total_pages"`
+}
+
+func main() {
+	path, err := tilde.Expand("~/.aws/credentials")
+	if err != nil {
+		log.Fatal("tilde.Expand: ", err)
+		return
+	}
+
+	cfg, err := ini.InsensitiveLoad(path)
+	if err != nil {
+		log.Fatal("ini.InsensitiveLoad: ", err)
+		return
+	}
+	mfa_serial := cfg.Section("default").Key("mfa_serial").String()
+	url := cfg.Section("default").Key("beagle_url").String()
+	pass := cfg.Section("default").Key("beagle_pass").String()
+
+	fmt.Printf("Locksmith GO\n")
+
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		log.Fatal("http.NewRequest: ", err)
+		return
+	}
+
+	req.SetBasicAuth("n/a", pass)
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		log.Fatal("client.Do: ", err)
+		return
+	}
+
+	// Callers should close resp.Body
+	// when done reading from it
+	// Defer the closing of the body
+	defer resp.Body.Close()
+
+	// Fill the record with the data from the JSON
+	var record Bookmarks
+
+	// Use json.Decode for reading streams of JSON data
+	if err := json.NewDecoder(resp.Body).Decode(&record); err != nil {
+		log.Println(err)
+	}
+
+	sort.Slice(record.Bookmarks, func(i, j int) bool {
+		return record.Bookmarks[i].Name < record.Bookmarks[j].Name
+	})
+
+	templates := &promptui.SelectTemplates{
+		Label:    "{{ . }}:",
+		Active:   "▸ {{ .AccountNumber | red }}: {{ .Name | yellow }}",
+		Inactive: "  {{ .AccountNumber | red | faint }}: {{ .Name | yellow | faint }}",
+		Selected: "{{ .AccountNumber | red }}: {{ .Name | yellow }}",
+		// 		Details: `
+		// --------- Account ----------
+		// {{ "Name:" | faint }}	{{ .Name }}
+		// {{ "AccountNumber:" | faint }}	{{ .AccountNumber }}
+		// {{ "RoleName:" | faint }}	{{ .RoleName }}`,
+	}
+
+	searcher := func(input string, index int) bool {
+		bookmark := record.Bookmarks[index]
+		name := strings.Replace(strings.ToLower(bookmark.Name), " ", "", -1)
+		input = strings.Replace(strings.ToLower(input), " ", "", -1)
+		name += bookmark.AccountNumber
+
+		return strings.Contains(name, input)
+	}
+
+	prompt := promptui.Select{
+		Label:     "AWS Account",
+		Items:     record.Bookmarks,
+		Templates: templates,
+		Size:      10,
+		Searcher:  searcher,
+	}
+
+	result, _, err := prompt.Run()
+
+	if err != nil {
+		fmt.Printf("Prompt failed %v\n", err)
+		return
+	}
+
+	validate := func(input string) error {
+		match, err := regexp.MatchString("^[0-9]{6}$", input)
+		if err != nil {
+			return err
+		}
+
+		if !match {
+			return errors.New("Token must be 6 digits")
+		}
+
+		return nil
+	}
+
+	mfaPrompt := promptui.Prompt{
+		Label:    "MFA Token",
+		Validate: validate,
+	}
+
+	token, err := mfaPrompt.Run()
+	if err != nil {
+		fmt.Printf("Prompt failed %v\n", err)
+		return
+	}
+
+	svc := sts.New(session.New())
+	input := &sts.AssumeRoleInput{
+		DurationSeconds: aws.Int64(3600),
+		RoleArn: aws.String(fmt.Sprintf(
+			"arn:aws:iam::%s:role/%s",
+			record.Bookmarks[result].AccountNumber,
+			record.Bookmarks[result].RoleName)),
+		RoleSessionName: aws.String("AssumeRoleSession"),
+		SerialNumber:    aws.String(mfa_serial),
+		TokenCode:       aws.String(token),
+	}
+
+	assumedRole, err := svc.AssumeRole(input)
+	if err != nil {
+		if aerr, ok := err.(awserr.Error); ok {
+			switch aerr.Code() {
+			case sts.ErrCodeMalformedPolicyDocumentException:
+				fmt.Println(sts.ErrCodeMalformedPolicyDocumentException, aerr.Error())
+			case sts.ErrCodePackedPolicyTooLargeException:
+				fmt.Println(sts.ErrCodePackedPolicyTooLargeException, aerr.Error())
+			case sts.ErrCodeRegionDisabledException:
+				fmt.Println(sts.ErrCodeRegionDisabledException, aerr.Error())
+			default:
+				fmt.Println(aerr.Error())
+			}
+		} else {
+			// Print the error, cast err to awserr.Error to get the Code and
+			// Message from an error.
+			fmt.Println(err.Error())
+		}
+		return
+	}
+
+	shell := user_shell.GetUserShell()
+	cmd := exec.Command(shell, "-l")
+	cmd.Env = append(os.Environ(),
+		fmt.Sprintf("AWS_ACCESS_KEY_ID=%s", aws.StringValue(assumedRole.Credentials.AccessKeyId)),
+		fmt.Sprintf("AWS_ASSUMED_ROLE_ARN=%s", aws.StringValue(input.RoleArn)),
+		fmt.Sprintf("AWS_SECRET_ACCESS_KEY=%s", aws.StringValue(assumedRole.Credentials.SecretAccessKey)),
+		fmt.Sprintf("AWS_SECURITY_TOKEN=%s", aws.StringValue(assumedRole.Credentials.SessionToken)),
+		fmt.Sprintf("AWS_SESSION_ACCOUNT_ID=%s", record.Bookmarks[result].AccountNumber),
+		fmt.Sprintf("AWS_SESSION_ACCOUNT_NAME=%s", record.Bookmarks[result].Name),
+		fmt.Sprintf("AWS_SESSION_EXPIRES=%d", aws.TimeValue(assumedRole.Credentials.Expiration).Unix()),
+		fmt.Sprintf("AWS_SESSION_TOKEN=%s", aws.StringValue(assumedRole.Credentials.SessionToken)),
+		fmt.Sprintf("AWS_SESSION_USER_ARN=%s", aws.StringValue(assumedRole.AssumedRoleUser.Arn)),
+		fmt.Sprintf("AWS_SESSION_USER_ID=%s", aws.StringValue(assumedRole.AssumedRoleUser.AssumedRoleId)),
+	)
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	cmdStartErr := cmd.Start()
+	if cmdStartErr != nil {
+		log.Fatal(cmdStartErr)
+	}
+	cmd.Wait()
+}