From 33eed42e17fda8b1a66feaeafcd57efccff26c11 Mon Sep 17 00:00:00 2001
From: Cody Rose <rosecodym@gmail.com>
Date: Fri, 25 Aug 2023 11:30:08 -0400
Subject: [PATCH] Test S3 role assumption (#1655)

This PR adds a test of the S3 role assumption functionality. It currently only tests role assumption within a single account.
---
 pkg/sources/s3/s3_test.go | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/pkg/sources/s3/s3_test.go b/pkg/sources/s3/s3_test.go
index f8f039ecd..c4e3fc75f 100644
--- a/pkg/sources/s3/s3_test.go
+++ b/pkg/sources/s3/s3_test.go
@@ -10,13 +10,12 @@ import (
 
 	"github.com/kylelemons/godebug/pretty"
 	"github.com/stretchr/testify/assert"
-	"google.golang.org/protobuf/types/known/anypb"
-
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/credentialspb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
+	"google.golang.org/protobuf/types/known/anypb"
 )
 
 func TestSource_Chunks(t *testing.T) {
@@ -35,6 +34,7 @@ func TestSource_Chunks(t *testing.T) {
 		name       string
 		verify     bool
 		connection *sourcespb.S3
+		setEnv     map[string]string
 	}
 	tests := []struct {
 		name          string
@@ -58,6 +58,23 @@ func TestSource_Chunks(t *testing.T) {
 			wantErr:       false,
 			wantChunkData: `W2RlZmF1bHRdCmF3c19hY2Nlc3Nfa2V5X2lkID0gQUtJQTM1T0hYMkRTT1pHNjQ3TkgKYXdzX3NlY3JldF9hY2Nlc3Nfa2V5ID0gUXk5OVMrWkIvQ1dsRk50eFBBaWQ3Z0d6dnNyWGhCQjd1ckFDQUxwWgpvdXRwdXQgPSBqc29uCnJlZ2lvbiA9IHVzLWVhc3QtMg==`,
 		},
+		{
+			name: "gets chunks after assuming role",
+			// This test will attempt to scan every bucket in the account, but the role policy blocks access to every
+			// bucket except the one we want. This (expected behavior) causes errors in the test log output, but these
+			// errors shouldn't actually cause test failures.
+			init: init{
+				connection: &sourcespb.S3{
+					Roles: []string{"arn:aws:iam::619888638459:role/s3-test-assume-role"},
+				},
+				setEnv: map[string]string{
+					"AWS_ACCESS_KEY_ID":     s3key,
+					"AWS_SECRET_ACCESS_KEY": s3secret,
+				},
+			},
+			wantErr:       false,
+			wantChunkData: `W2RlZmF1bHRdCmF3c19zZWNyZXRfYWNjZXNzX2tleSA9IFF5OTlTK1pCL0NXbEZOdHhQQWlkN2dHenZzclhoQkI3dXJBQ0FMcFoKYXdzX2FjY2Vzc19rZXlfaWQgPSBBS0lBMzVPSFgyRFNPWkc2NDdOSApvdXRwdXQgPSBqc29uCnJlZ2lvbiA9IHVzLWVhc3QtMg==`,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -65,6 +82,10 @@ func TestSource_Chunks(t *testing.T) {
 			var cancelOnce sync.Once
 			defer cancelOnce.Do(cancel)
 
+			for k, v := range tt.init.setEnv {
+				t.Setenv(k, v)
+			}
+
 			s := Source{}
 			conn, err := anypb.New(tt.init.connection)
 			if err != nil {