<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"io/ioutil"
	"os"
	"os/user"
	"strconv"
	"strings"
	"syscall"
	"time"
)

func main() {
	// Subcommands
	authCommand := flag.NewFlagSet("auth", flag.ExitOnError)
	switchCommand := flag.NewFlagSet("switch", flag.ExitOnError)

	// Auth Options
	profilePtr := authCommand.String("profile", "", "The amazon credentials profile")
	accountNumberPtr := authCommand.String("account", "", "The amazon account of your MFA device")
	userNamePtr := authCommand.String("username", "", "The amazon username associated with your MFA device")
	tokenCodePtr := authCommand.String("token", "", "The OTP token to use")
	savePtr := authCommand.Bool("save", false, "Save the supplied details")

	// Switch Options
	profilePtr = switchCommand.String("profile", "", "The amazon credentials profile")
	roleAccountNumberPtr := switchCommand.String("account", "", "The amazon account of your role")
	roleNamePtr := switchCommand.String("role", "", "The amazon role name associated to assume")
	saveProfilePtr := switchCommand.Bool("save", false, "Save the supplied details")

	// Verify that a subcommand has been provided
	// os.Arg[0] is the main command
	// os.Arg[1] will be the subcommand
	if len(os.Args) < 2 {
		fmt.Println("auth, switch or clear subcommand is required")
		os.Exit(1)
	}

	profile := "default"
	envProfile := os.Getenv("PORTRAY_PROFILE")
	fileName := ""
	roleFileName := ""

	// Get our session file
	usr, err := user.Current()
	checkError(err)
	portrayConfigFileName := usr.HomeDir + "/.portray-config.json"
	config := getPortrayConfigFromFile(portrayConfigFileName)

	// Switch on the subcommand
	// Parse the flags for appropriate FlagSet
	// FlagSet.Parse() requires a set of arguments to parse as input
	// os.Args[2:] will be all arguments starting after the subcommand at os.Args[1]
	switch os.Args[1] {
	case "auth":
		authCommand.Parse(os.Args[2:])
		if *profilePtr != "" {
			profile = *profilePtr
			os.Setenv("PORTRAY_PROFILE", profile)
		} else if envProfile != "" {
			profile = envProfile
		} else if config.DefaultProfile != "" {
			profile = config.DefaultProfile
		}
		fileName = usr.HomeDir + "/.aws/portray-session-" + profile + ".json"
	case "switch":
		switchCommand.Parse(os.Args[2:])
		if *profilePtr != "" {
			profile = *profilePtr
			os.Setenv("PORTRAY_PROFILE", profile)
		} else if envProfile != "" {
			profile = envProfile
		} else if config.DefaultProfile != "" {
			profile = config.DefaultProfile
		}
		fileName = usr.HomeDir + "/.aws/portray-session-" + profile + ".json"
		roleFileName = usr.HomeDir + "/.aws/portray-role-session-" + *roleAccountNumberPtr + "_" + *roleNamePtr + ".json"
	default:
		flag.PrintDefaults()
		os.Exit(1)
	}

	if authCommand.Parsed() {
		awsCreds := getCredsFromFile(fileName)
		if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
			if *tokenCodePtr != "" {
				awsCreds = getNewSession(*accountNumberPtr, *userNamePtr, *tokenCodePtr)
			} else if config.AccountNumber != "" {
				reader := bufio.NewReader(os.Stdin)
				fmt.Print("Enter token: ")
				token, _ := reader.ReadString('\n')
				token = strings.TrimSpace(token)
				awsCreds = getNewSession(config.AccountNumber, config.UserName, token)
			} else {
				fmt.Println("You need a valid session!")
				os.Exit(1)
			}
			writeSessionFile(awsCreds, fileName)
		}
		if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
			fmt.Println("You need a valid session!")
			os.Exit(1)
		}

		accountNumber := awsCreds.AccountNumber

		sessionToEnvVars(awsCreds, accountNumber, "", profile)
		if *savePtr == true {
			if *profilePtr != "" {
				config.DefaultProfile = *profilePtr
			}
			if *accountNumberPtr != "" {
				config.AccountNumber = *accountNumberPtr
			}

			if *userNamePtr != "" {
				config.UserName = *userNamePtr
			}
			writePortrayConfigToFile(portrayConfigFileName, config)
		}
		startShell(accountNumber)
	} else if switchCommand.Parsed() {
		profileIdx := -1

		if *profilePtr != "" {
			for i := range config.Profiles {
				if config.Profiles[i].Name == *profilePtr {
					profileIdx = i
					fmt.Println("Found Profile")
					roleFileName = usr.HomeDir + "/.aws/portray-role-session-" + config.Profiles[i].AccountNumber + "_" + config.Profiles[i].RoleName + ".json"
					break
				}
			}
		}

		awsRoleCreds := getCredsFromFile(roleFileName)
		if awsRoleCreds.SessionToken == "" || !validateSession(awsRoleCreds) {
			fileName = usr.HomeDir + "/.aws/portray-session-" + envProfile + ".json"
			awsCreds := getCredsFromFile(fileName)
			if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
				if config.AccountNumber != "" {
					fmt.Printf("Attempt auth using config...")
					reader := bufio.NewReader(os.Stdin)
					fmt.Print("Enter token: ")
					token, _ := reader.ReadString('\n')
					token = strings.TrimSpace(token)
					awsCreds = getNewSession(config.AccountNumber, config.UserName, token)
					writeSessionFile(awsCreds, fileName)
				} else {
					fmt.Println("You need a valid session!")
					os.Exit(1)
				}
			}
			if *roleAccountNumberPtr == "" && profileIdx == -1 {
				fmt.Println("You must have a profile configured or supply a role account #")
				os.Exit(1)
			}
			if *roleAccountNumberPtr != "" {
				awsRoleCreds = getNewRoleSession(*roleAccountNumberPtr, *roleNamePtr, *usr)
			} else {
				awsRoleCreds = getNewRoleSession(
					config.Profiles[profileIdx].AccountNumber,
					config.Profiles[profileIdx].RoleName,
					*usr,
				)
			}
			writeSessionFile(awsRoleCreds, roleFileName)
		}

		roleAccountNumber := awsRoleCreds.AccountNumber
		roleName := awsRoleCreds.RoleName

		sessionToEnvVars(awsRoleCreds, roleAccountNumber, roleName, profile)
		if *saveProfilePtr == true {
			if profileIdx >= 0 {
				config.Profiles[profileIdx].RoleName = *roleNamePtr
				config.Profiles[profileIdx].AccountNumber = *roleAccountNumberPtr
				config.Profiles[profileIdx].Name = *profilePtr
			} else if *profilePtr != "" {
				profile := PortrayProfile{*profilePtr, *roleAccountNumberPtr, *roleNamePtr}
				config.Profiles = append(config.Profiles, profile)
			}
			writePortrayConfigToFile(portrayConfigFileName, config)
		}
		sessionName := ""
		if *roleNamePtr == "" {
			sessionName := roleAccountNumber + "-" + roleName
		} else {
			sessionName := roleAccountNumber + "-" + *roleNamePtr
		}
		startShell(sessionName)
	}
}

// AwsCreds represents a set of AWS credentials
type AwsCreds struct {
	AccessKeyID     string
	Expiration      int64
	SecretAccessKey string
	SessionToken    string
	AccountNumber   string
	RoleName        string
}

// PortrayProfile represents a valid AWS assumable Role
type PortrayProfile struct {
	Name          string
	AccountNumber string<commit_msg>Fixing go 1.9 sessionName redefine<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"io/ioutil"
	"os"
	"os/user"
	"strconv"
	"strings"
	"syscall"
	"time"
)

func main() {
	// Subcommands
	authCommand := flag.NewFlagSet("auth", flag.ExitOnError)
	switchCommand := flag.NewFlagSet("switch", flag.ExitOnError)

	// Auth Options
	profilePtr := authCommand.String("profile", "", "The amazon credentials profile")
	accountNumberPtr := authCommand.String("account", "", "The amazon account of your MFA device")
	userNamePtr := authCommand.String("username", "", "The amazon username associated with your MFA device")
	tokenCodePtr := authCommand.String("token", "", "The OTP token to use")
	savePtr := authCommand.Bool("save", false, "Save the supplied details")

	// Switch Options
	profilePtr = switchCommand.String("profile", "", "The amazon credentials profile")
	roleAccountNumberPtr := switchCommand.String("account", "", "The amazon account of your role")
	roleNamePtr := switchCommand.String("role", "", "The amazon role name associated to assume")
	saveProfilePtr := switchCommand.Bool("save", false, "Save the supplied details")

	// Verify that a subcommand has been provided
	// os.Arg[0] is the main command
	// os.Arg[1] will be the subcommand
	if len(os.Args) < 2 {
		fmt.Println("auth, switch or clear subcommand is required")
		os.Exit(1)
	}

	profile := "default"
	envProfile := os.Getenv("PORTRAY_PROFILE")
	fileName := ""
	roleFileName := ""

	// Get our session file
	usr, err := user.Current()
	checkError(err)
	portrayConfigFileName := usr.HomeDir + "/.portray-config.json"
	config := getPortrayConfigFromFile(portrayConfigFileName)

	// Switch on the subcommand
	// Parse the flags for appropriate FlagSet
	// FlagSet.Parse() requires a set of arguments to parse as input
	// os.Args[2:] will be all arguments starting after the subcommand at os.Args[1]
	switch os.Args[1] {
	case "auth":
		authCommand.Parse(os.Args[2:])
		if *profilePtr != "" {
			profile = *profilePtr
			os.Setenv("PORTRAY_PROFILE", profile)
		} else if envProfile != "" {
			profile = envProfile
		} else if config.DefaultProfile != "" {
			profile = config.DefaultProfile
		}
		fileName = usr.HomeDir + "/.aws/portray-session-" + profile + ".json"
	case "switch":
		switchCommand.Parse(os.Args[2:])
		if *profilePtr != "" {
			profile = *profilePtr
			os.Setenv("PORTRAY_PROFILE", profile)
		} else if envProfile != "" {
			profile = envProfile
		} else if config.DefaultProfile != "" {
			profile = config.DefaultProfile
		}
		fileName = usr.HomeDir + "/.aws/portray-session-" + profile + ".json"
		roleFileName = usr.HomeDir + "/.aws/portray-role-session-" + *roleAccountNumberPtr + "_" + *roleNamePtr + ".json"
	default:
		flag.PrintDefaults()
		os.Exit(1)
	}

	if authCommand.Parsed() {
		awsCreds := getCredsFromFile(fileName)
		if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
			if *tokenCodePtr != "" {
				awsCreds = getNewSession(*accountNumberPtr, *userNamePtr, *tokenCodePtr)
			} else if config.AccountNumber != "" {
				reader := bufio.NewReader(os.Stdin)
				fmt.Print("Enter token: ")
				token, _ := reader.ReadString('\n')
				token = strings.TrimSpace(token)
				awsCreds = getNewSession(config.AccountNumber, config.UserName, token)
			} else {
				fmt.Println("You need a valid session!")
				os.Exit(1)
			}
			writeSessionFile(awsCreds, fileName)
		}
		if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
			fmt.Println("You need a valid session!")
			os.Exit(1)
		}

		accountNumber := awsCreds.AccountNumber

		sessionToEnvVars(awsCreds, accountNumber, "", profile)
		if *savePtr == true {
			if *profilePtr != "" {
				config.DefaultProfile = *profilePtr
			}
			if *accountNumberPtr != "" {
				config.AccountNumber = *accountNumberPtr
			}

			if *userNamePtr != "" {
				config.UserName = *userNamePtr
			}
			writePortrayConfigToFile(portrayConfigFileName, config)
		}
		startShell(accountNumber)
	} else if switchCommand.Parsed() {
		profileIdx := -1

		if *profilePtr != "" {
			for i := range config.Profiles {
				if config.Profiles[i].Name == *profilePtr {
					profileIdx = i
					fmt.Println("Found Profile")
					roleFileName = usr.HomeDir + "/.aws/portray-role-session-" + config.Profiles[i].AccountNumber + "_" + config.Profiles[i].RoleName + ".json"
					break
				}
			}
		}

		awsRoleCreds := getCredsFromFile(roleFileName)
		if awsRoleCreds.SessionToken == "" || !validateSession(awsRoleCreds) {
			fileName = usr.HomeDir + "/.aws/portray-session-" + envProfile + ".json"
			awsCreds := getCredsFromFile(fileName)
			if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
				if config.AccountNumber != "" {
					fmt.Printf("Attempt auth using config...")
					reader := bufio.NewReader(os.Stdin)
					fmt.Print("Enter token: ")
					token, _ := reader.ReadString('\n')
					token = strings.TrimSpace(token)
					awsCreds = getNewSession(config.AccountNumber, config.UserName, token)
					writeSessionFile(awsCreds, fileName)
				} else {
					fmt.Println("You need a valid session!")
					os.Exit(1)
				}
			}
			if *roleAccountNumberPtr == "" && profileIdx == -1 {
				fmt.Println("You must have a profile configured or supply a role account #")
				os.Exit(1)
			}
			if *roleAccountNumberPtr != "" {
				awsRoleCreds = getNewRoleSession(*roleAccountNumberPtr, *roleNamePtr, *usr)
			} else {
				awsRoleCreds = getNewRoleSession(
					config.Profiles[profileIdx].AccountNumber,
					config.Profiles[profileIdx].RoleName,
					*usr,
				)
			}
			writeSessionFile(awsRoleCreds, roleFileName)
		}

		roleAccountNumber := awsRoleCreds.AccountNumber
		roleName := awsRoleCreds.RoleName

		sessionToEnvVars(awsRoleCreds, roleAccountNumber, roleName, profile)
		if *saveProfilePtr == true {
			if profileIdx >= 0 {
				config.Profiles[profileIdx].RoleName = *roleNamePtr
				config.Profiles[profileIdx].AccountNumber = *roleAccountNumberPtr
				config.Profiles[profileIdx].Name = *profilePtr
			} else if *profilePtr != "" {
				profile := PortrayProfile{*profilePtr, *roleAccountNumberPtr, *roleNamePtr}
				config.Profiles = append(config.Profiles, profile)
			}
			writePortrayConfigToFile(portrayConfigFileName, config)
		}
		sessionName := ""
		if *roleNamePtr == "" {
			sessionName = roleAccountNumber + "-" + roleName
		} else {
			sessionName = roleAccountNumber + "-" + *roleNamePtr
		}
		startShell(sessionName)
	}
}

// AwsCreds represents a set of AWS credentials
type AwsCreds struct {
	AccessKeyID     string
	Expiration      int64
	SecretAccessKey string
	SessionToken    string
	AccountNumber   string
	RoleName        string
}

// PortrayProfile represents a valid AWS assumable Role
type PortrayProfile struct {
	Name          string
	AccountNumber string
{'fraction_non_alphanumeric': 0.1333187868208597, 'fraction_numerical': 0.002909302494726889, 'mean_word_length': 3.8347398030942337, 'pattern_counts': {'":': 4, '<': 5, '<?xml version=': 0, '>': 5, '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': '5666533', 'n_tokens_mistral': 4989, 'n_tokens_neox': 4435, 'n_words': 965}