From fae7316118592ba6f9481d9389669fba0ac575e0 Mon Sep 17 00:00:00 2001
From: Edward Viaene <ward@in4it.io>
Date: Wed, 6 Apr 2022 11:44:36 -0400
Subject: [PATCH] duration seconds credential

---
 README.md                   |  1 +
 cmd/login.go                |  6 +++++-
 cmd/root.go                 | 12 +++++++-----
 internal/aws/assume-role.go |  4 ++--
 4 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index dbd62ef..63ee677 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,7 @@ onelogin:
   clientID: clientID of API credential with "Authentication only"
   clientSecret: client Secret of API credential
   accountName: onelogin account name
+  durationSeconds: 28800 # duration of the credentials in seconds (or remove for the default of 3600)
 accounts:
   - name: myapp-prod
     appID: onelogin app id (e.g. 123456)
diff --git a/cmd/login.go b/cmd/login.go
index 52400f5..1f52017 100644
--- a/cmd/login.go
+++ b/cmd/login.go
@@ -97,7 +97,11 @@ var loginCmd = &cobra.Command{
 		//AssumeRole With SAML on AWS
 		accountID := config.Accounts[*account].AccountID
 		profileName := config.Accounts[*account].ProfileName
-		result, err := intAWS.AssumeRoleWithSAML(accountID, config.Roles[*role], assertionPayload)
+		durationSeconds := config.Accounts[*account].DurationSeconds
+		if durationSeconds == 0 {
+			durationSeconds = 3600
+		}
+		result, err := intAWS.AssumeRoleWithSAML(accountID, config.Roles[*role], assertionPayload, durationSeconds)
 		if err != nil {
 			fmt.Println(err)
 			os.Exit(1)
diff --git a/cmd/root.go b/cmd/root.go
index f477db0..bf610ae 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -1,9 +1,10 @@
 package cmd
 
 import (
+	"log"
+
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-	"log"
 )
 
 type Config struct {
@@ -19,10 +20,11 @@ type OneLoginConf struct {
 	AccountName  string `yaml:"onelogin-account"`
 }
 type Account struct {
-	Name        string `yaml:"name"`
-	AppID       string `yaml:"appID"`
-	AccountID   string `yaml:"accountID"`
-	ProfileName string `yaml:"profileName"`
+	Name            string `yaml:"name"`
+	AppID           string `yaml:"appID"`
+	AccountID       string `yaml:"accountID"`
+	ProfileName     string `yaml:"profileName"`
+	DurationSeconds int64  `yaml:"durationSeconds"`
 }
 
 var config Config
diff --git a/internal/aws/assume-role.go b/internal/aws/assume-role.go
index 608c1dd..ecdbf9b 100644
--- a/internal/aws/assume-role.go
+++ b/internal/aws/assume-role.go
@@ -6,7 +6,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/sts"
 )
 
-func AssumeRoleWithSAML(accountID string, role string, assertionPayload string) (*sts.AssumeRoleWithSAMLOutput, error) {
+func AssumeRoleWithSAML(accountID string, role string, assertionPayload string, durationSeconds int64) (*sts.AssumeRoleWithSAMLOutput, error) {
 	sess, err := session.NewSession(&aws.Config{
 		Region: aws.String("us-east-1"),
 	})
@@ -19,7 +19,7 @@ func AssumeRoleWithSAML(accountID string, role string, assertionPayload string)
 	roleToAssumeArn := "arn:aws:iam::" + accountID + ":role/" + role
 	result, err := svc.AssumeRoleWithSAML(&sts.AssumeRoleWithSAMLInput{
 		RoleArn:         &roleToAssumeArn,
-		DurationSeconds: aws.Int64(3600),
+		DurationSeconds: aws.Int64(durationSeconds),
 		PrincipalArn:    aws.String("arn:aws:iam::" + accountID + ":saml-provider/" + role),
 		SAMLAssertion:   aws.String(assertionPayload),
 	})