Specified by option

youyo · May 18, 2018 · 5794146 · 5794146
1 parent 001ce44
commit 5794146
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 13 deletions.
diff --git a/awslogin.go b/awslogin.go
@@ -35,15 +35,15 @@ func NewSession(sourceProfile string) (s *session.Session, err error) {
 	return
 }
 
-func NewCredentials(sess *session.Session, arn, roleSessionName, mfaSerial string) (creds credentials.Value, err error) {
-	assumeRoleProvider := buildAssumeRoleProvider(roleSessionName, mfaSerial)
+func NewCredentials(sess *session.Session, arn, roleSessionName, mfaSerial string, durationSeconds int) (creds credentials.Value, err error) {
+	assumeRoleProvider := buildAssumeRoleProvider(roleSessionName, mfaSerial, durationSeconds)
 	creds, err = stscreds.NewCredentials(sess, arn, assumeRoleProvider).Get()
 	return
 }
 
-func buildAssumeRoleProvider(roleSessionName, mfaSerial string) (f func(p *stscreds.AssumeRoleProvider)) {
+func buildAssumeRoleProvider(roleSessionName, mfaSerial string, durationSeconds int) (f func(p *stscreds.AssumeRoleProvider)) {
 	f = func(p *stscreds.AssumeRoleProvider) {
-		p.Duration = time.Duration(60*60*12) * time.Second
+		p.Duration = time.Duration(durationSeconds) * time.Second
 		p.RoleSessionName = roleSessionName
 		if mfaSerial != "" {
 			p.SerialNumber = aws.String(mfaSerial)

diff --git a/awslogin_test.go b/awslogin_test.go
@@ -23,6 +23,7 @@ func TestBuildAssumeRoleProvider(t *testing.T) {
 	p := buildAssumeRoleProvider(
 		"test.xxxxxxxxxxxx@awslogin",
 		"arn:aws:iam::123456789012:mfa/jonsmith",
+		3600,
 	)
 	var expected func(p *stscreds.AssumeRoleProvider)
 	if reflect.TypeOf(p) != reflect.TypeOf(expected) {

diff --git a/cmd/awslogin/cmd/root.go b/cmd/awslogin/cmd/root.go
@@ -15,14 +15,15 @@ import (
 )
 
 var (
-	Name        string
-	Version     string
-	CommitHash  string
-	BuildTime   string
-	GoVersion   string
-	app         string
-	profile     string
-	readFromEnv bool
+	Name            string
+	Version         string
+	CommitHash      string
+	BuildTime       string
+	GoVersion       string
+	app             string
+	profile         string
+	readFromEnv     bool
+	durationSeconds int
 )
 
 const (
@@ -63,7 +64,7 @@ func execRoot() (err error) {
 		return
 	}
 
-	creds, err := awslogin.NewCredentials(sess, cfg.ARN, cfg.RoleSessionName, cfg.MfaSerial)
+	creds, err := awslogin.NewCredentials(sess, cfg.ARN, cfg.RoleSessionName, cfg.MfaSerial, durationSeconds)
 	if err != nil {
 		return
 	}
@@ -145,6 +146,7 @@ func init() {
 	RootCmd.Flags().StringVarP(&app, "app", "a", "", "Opens with the specified application.")
 	RootCmd.Flags().StringVarP(&profile, "profile", "p", "", "Use a specific profile.")
 	RootCmd.Flags().BoolVarP(&readFromEnv, "read-from-env", "e", false, "Use a specific profile read from the environment. [$AWS_PROFILE]")
+	RootCmd.Flags().IntVarP(&durationSeconds, "duration-seconds", "d", 3600, "Request a session duration seconds. 900 - 43200")
 }
 
 func initConfig() {}