<commit_before>package cmd

import (
	"fmt"
	"os"
	"regexp"
	"sort"

	"github.com/akerl/voyager/cartogram"

	"github.com/dixonwille/wmenu"
	"github.com/spf13/cobra"
)

const (
	accountRegexString = `(\d+)(/(\w+))?`
)

var accountRegex = regexp.MustCompile(accountRegexString)

var travelCmd = &cobra.Command{
	Use:   "travel",
	Short: "Resolve creds for a AWS account",
	RunE:  travelRunner,
}

func init() {
	rootCmd.AddCommand(travelCmd)
}

func travelRunner(cmd *cobra.Command, args []string) error {
	cp := cartogram.Pack{}
	if err := cp.Load(); err != nil {
		return err
	}

	targetAccount, targetRole, err := findAccount(cp, args)
	if err != nil {
		return err
	}

	if targetRole == "" {
		roleNames := []string{}
		for k := range targetAccount.Roles {
			roleNames = append(roleNames, k)
		}
		sort.Strings(roleNames)
		targetRole, err = pickFromList("Desired Role:", roleNames, "")
		if err != nil {
			return err
		}
	}

	fmt.Printf("Account: %s\nRole: %s\n", targetAccount, targetRole)
	return nil
}

func findAccount(cp cartogram.Pack, args []string) (cartogram.Account, string, error) {
	var targetAccount cartogram.Account
	var targetRole string
	var err error
	var found bool

	found, targetAccount, targetRole, err = findDirectAccount(cp, args)
	if err != nil {
		return targetAccount, "", err
	}
	if found {
		return targetAccount, targetRole, nil
	}

	found, targetAccount, targetRole, err = findMatchAccount(cp, args)
	if found {
		return targetAccount, targetRole, nil
	}

	return targetAccount, targetRole, fmt.Errorf("Unable to locate an account with provided info")
}

func findDirectAccount(cp cartogram.Pack, args []string) (bool, cartogram.Account, string, error) {
	var account cartogram.Account
	if len(args) != 1 {
		return false, account, "", nil
	}
	accountMatch := accountRegex.FindStringSubmatch(args[0])
	if len(accountMatch) == 0 {<commit_msg>update with working janky travel<commit_after>package cmd

import (
	"fmt"
	"os"
	"regexp"
	"sort"

	"github.com/akerl/voyager/cartogram"

	speculate "github.com/akerl/speculate/utils"
	"github.com/dixonwille/wmenu"
	"github.com/spf13/cobra"
)

const (
	accountRegexString = `(\d+)(/(\w+))?`
)

var accountRegex = regexp.MustCompile(accountRegexString)

var travelCmd = &cobra.Command{
	Use:   "travel",
	Short: "Resolve creds for a AWS account",
	RunE:  travelRunner,
}

func init() {
	rootCmd.AddCommand(travelCmd)
	travelCmd.Flags().StringP("role", "r", "", "Choose role to use")
}

type hop struct {
	Profile string
	Account string
	Role    string
	Mfa     bool
}

func travelRunner(cmd *cobra.Command, args []string) error {
	flags := cmd.Flags()

	cp := cartogram.Pack{}
	if err := cp.Load(); err != nil {
		return err
	}

	targetAccount, targetRole, err := findAccount(cp, args)
	if err != nil {
		return err
	}

	if targetRole == "" {
		// TODO: check if role exists
		targetRole, err = flags.GetString("role")
		if err != nil {
			return err
		}
	}

	if targetRole == "" {
		roleNames := []string{}
		for k := range targetAccount.Roles {
			roleNames = append(roleNames, k)
		}
		sort.Strings(roleNames)
		// TODO: don't ask if there's only 1 option
		targetRole, err = pickFromList("Desired Role:", roleNames, "")
		if err != nil {
			return err
		}
	}

	stack := []hop{}
	if err := parseHops(&stack, cp, targetAccount, targetRole); err != nil {
		return err
	}
	for i, j := 0, len(stack)-1; i < j; i, j = i+1, j-1 {
		stack[i], stack[j] = stack[j], stack[i]
	}

	for _, thisHop := range stack {
		// TODO: pop the profile off the top pre-loop
		if thisHop.Profile != "" {
			os.Setenv("AWS_PROFILE", thisHop.Profile)
		} else {
			assumption := speculate.Assumption{
				RoleName:  thisHop.Role,
				AccountID: thisHop.Account,
			}
			// TODO: make speculate default to lifetime 3600 if not set
			assumption.Lifetime.LifetimeInt = 3600
			assumption.Mfa.UseMfa = thisHop.Mfa
			creds, err := assumption.Execute()
			if err != nil {
				return err
			}
			// TODO: make speculate take creds as an assumption field
			// TODO: make translations visible externally
			os.Setenv("AWS_ACCESS_KEY_ID", creds.AccessKey)
			os.Setenv("AWS_SECRET_ACCESS_KEY", creds.SecretKey)
			os.Setenv("AWS_SESSION_TOKEN", creds.SessionToken)
			os.Setenv("AWS_SECURITY_TOKEN", creds.SessionToken)
		}
	}

	//TODO: Avoid double-initing these creds
	creds := speculate.Creds{}
	if err := creds.NewFromEnv(); err != nil {
		return err
	}
	for _, line := range creds.ToEnvVars() {
		fmt.Println(line)
	}
	url, err := creds.ToConsoleURL()
	if err != nil {
		return err
	}
	fmt.Printf("# %s\n", url)

	return nil
}

func parseHops(stack *[]hop, cp cartogram.Pack, a cartogram.Account, r string) error {
	*stack = append(*stack, hop{Account: a.Account, Role: r, Mfa: a.Roles[r].Mfa})
	accountMatch := accountRegex.FindStringSubmatch(a.Source)
	if len(accountMatch) != 4 {
		*stack = append(*stack, hop{Profile: a.Source})
		return nil
	}
	sAccountID := accountMatch[1]
	sRole := accountMatch[3]
	found, sAccount := cp.Lookup(sAccountID)
	if !found {
		return fmt.Errorf("Failed to resolve hop for %s", sAccountID)
	}
	return parseHops(stack, cp, sAccount, sRole)
}

func findAccount(cp cartogram.Pack, args []string) (cartogram.Account, string, error) {
	var targetAccount cartogram.Account
	var targetRole string
	var err error
	var found bool

	found, targetAccount, targetRole, err = findDirectAccount(cp, args)
	if err != nil {
		return targetAccount, "", err
	}
	if found {
		return targetAccount, targetRole, nil
	}

	found, targetAccount, targetRole, err = findMatchAccount(cp, args)
	if found {
		return targetAccount, targetRole, nil
	}

	return targetAccount, targetRole, fmt.Errorf("Unable to locate an account with provided info")
}

func findDirectAccount(cp cartogram.Pack, args []string) (bool, cartogram.Account, string, error) {
	var account cartogram.Account
	if len(args) != 1 {
		return false, account, "", nil
	}
	accountMatch := accountRegex.FindStringSubmatch(args[0])
	if len(accountMatch) == 0 {
{'fraction_non_alphanumeric': 0.13421360521477518, 'fraction_numerical': 0.004345645996991476, 'mean_word_length': 4.011725293132328, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7097153', 'n_tokens_mistral': 2242, 'n_tokens_neox': 2093, 'n_words': 562}