From b9fb8626499fbc1191c70732cb625ff6528eecf0 Mon Sep 17 00:00:00 2001
From: Ompragash Viswanathan <ompragash@proton.me>
Date: Fri, 19 Apr 2024 17:49:13 +0530
Subject: [PATCH 1/2] Added PLUGIN_OIDC_TOKEN_ID support

---
 main.go   |  6 +++++
 plugin.go | 81 ++++++++++++++++++++++++++++++++++++-------------------
 2 files changed, 60 insertions(+), 27 deletions(-)

diff --git a/main.go b/main.go
index a1114f4..6cf56a8 100644
--- a/main.go
+++ b/main.go
@@ -141,6 +141,11 @@ func main() {
 			Usage:  "external ID to use when assuming role",
 			EnvVar: "PLUGIN_EXTERNAL_ID",
 		},
+		cli.StringFlag{
+			Name:   "oidc-token-id",
+			Usage:  "OIDC token for assuming role via web identity",
+			EnvVar: "PLUGIN_OIDC_TOKEN_ID",
+		},
 	}
 
 	if err := app.Run(os.Args); err != nil {
@@ -176,6 +181,7 @@ func run(c *cli.Context) error {
 		PathStyle:             c.Bool("path-style"),
 		DryRun:                c.Bool("dry-run"),
 		ExternalID:            c.String("external-id"),
+		IdToken: 			   c.String("oidc-token-id"),
 	}
 
 	return plugin.Exec()
diff --git a/plugin.go b/plugin.go
index d7b3387..4b461e7 100644
--- a/plugin.go
+++ b/plugin.go
@@ -98,6 +98,9 @@ type Plugin struct {
 
 	// set externalID for assume role
 	ExternalID string
+
+	// set OIDC ID Token to retrieve temporary credentials 
+	IdToken string
 }
 
 // Exec runs the plugin
@@ -431,31 +434,55 @@ func (p *Plugin) downloadS3Objects(client *s3.S3, sourceDir string) error {
 
 // createS3Client creates and returns an S3 client based on the plugin configuration
 func (p *Plugin) createS3Client() *s3.S3 {
-	conf := &aws.Config{
-		Region:           aws.String(p.Region),
-		Endpoint:         &p.Endpoint,
-		DisableSSL:       aws.Bool(strings.HasPrefix(p.Endpoint, "http://")),
-		S3ForcePathStyle: aws.Bool(p.PathStyle),
-	}
-
-	if p.Key != "" && p.Secret != "" {
-		conf.Credentials = credentials.NewStaticCredentials(p.Key, p.Secret, "")
-	} else if p.AssumeRole != "" {
-		conf.Credentials = assumeRole(p.AssumeRole, p.AssumeRoleSessionName, p.ExternalID)
-	} else {
-		log.Warn("AWS Key and/or Secret not provided (falling back to ec2 instance profile)")
-	}
-
-	sess, _ := session.NewSession(conf)
-	client := s3.New(sess)
-
-	if len(p.UserRoleArn) > 0 {
-		confRoleArn := aws.Config{
-			Region:      aws.String(p.Region),
-			Credentials: stscreds.NewCredentials(sess, p.UserRoleArn),
-		}
-		client = s3.New(sess, &confRoleArn)
-	}
-
-	return client
+    conf := &aws.Config{
+        Region:           aws.String(p.Region),
+        Endpoint:         &p.Endpoint,
+        DisableSSL:       aws.Bool(strings.HasPrefix(p.Endpoint, "http://")),
+        S3ForcePathStyle: aws.Bool(p.PathStyle),
+    }
+
+    sess, err := session.NewSession(conf)
+    if err != nil {
+        log.Fatalf("failed to create AWS session: %v", err)
+    }
+
+    if p.Key != "" && p.Secret != "" {
+        conf.Credentials = credentials.NewStaticCredentials(p.Key, p.Secret, "")
+    } else if p.IdToken != "" && p.AssumeRole != "" {
+        creds, err := assumeRoleWithWebIdentity(sess, p.AssumeRole, p.AssumeRoleSessionName, p.IdToken)
+        if err != nil {
+            log.Fatalf("failed to assume role with web identity: %v", err)
+        }
+        conf.Credentials = creds
+    } else if p.AssumeRole != "" {
+        conf.Credentials = assumeRole(p.AssumeRole, p.AssumeRoleSessionName, p.ExternalID)
+    } else {
+        log.Warn("AWS Key and/or Secret not provided (falling back to ec2 instance profile)")
+    }
+
+    client := s3.New(sess, conf)
+
+    if len(p.UserRoleArn) > 0 {
+        confRoleArn := aws.Config{
+            Region:      aws.String(p.Region),
+            Credentials: stscreds.NewCredentials(sess, p.UserRoleArn),
+        }
+        client = s3.New(sess, &confRoleArn)
+    }
+
+    return client
 }
+
+func assumeRoleWithWebIdentity(sess *session.Session, roleArn, roleSessionName, idToken string) (*credentials.Credentials, error) {
+    svc := sts.New(sess)
+    input := &sts.AssumeRoleWithWebIdentityInput{
+        RoleArn:          aws.String(roleArn),
+        RoleSessionName:  aws.String(roleSessionName),
+        WebIdentityToken: aws.String(idToken),
+    }
+    result, err := svc.AssumeRoleWithWebIdentity(input)
+    if err != nil {
+        log.Fatalf("failed to assume role with web identity: %v", err)
+    }
+    return credentials.NewStaticCredentials(*result.Credentials.AccessKeyId, *result.Credentials.SecretAccessKey, *result.Credentials.SessionToken), nil
+}
\ No newline at end of file

From 4d577dfa34ceae76074643ff7391dcf2806487d3 Mon Sep 17 00:00:00 2001
From: Ompragash Viswanathan <ompragash@proton.me>
Date: Tue, 23 Apr 2024 18:54:40 +0530
Subject: [PATCH 2/2] Trigger notification