Changes and fixes

Author: shiv-am0

go.mod

@@ -13,6 +13,7 @@ require (
 require (
 	github.com/cpuguy83/go-md2man/v2 v2.0.1 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/pkg/errors v0.9.1
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	golang.org/x/sync v0.6.0
 	golang.org/x/sys v0.1.0 // indirect

go.sum

@@ -15,6 +15,7 @@ github.com/joho/godotenv v1.4.0 h1:3l4+N6zfMWnkbPEXKng2o2/MR5mSwTrBih4ZEkkz1lg=
 github.com/joho/godotenv v1.4.0/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
 github.com/mattn/go-zglob v0.0.4 h1:LQi2iOm0/fGgu80AioIJ/1j9w9Oh+9DZ39J4VAGzHQM=
 github.com/mattn/go-zglob v0.0.4/go.mod h1:MxxjyoXXnMxfIpxTK2GAkw1w8glPsQILx3N5wrKakiY=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=

plugin.go

@@ -16,6 +16,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/aws/aws-sdk-go/service/sts"
 	"github.com/mattn/go-zglob"
+	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -101,11 +102,11 @@ type Plugin struct {
 
 // Exec runs the plugin
 func (p *Plugin) Exec() error {
-	// normalize the target URL
 	if p.Download {
 		p.Source = normalizePath(p.Source)
-	} else {
 		p.Target = normalizePath(p.Target)
+	} else {
+		p.Target = strings.TrimPrefix(p.Target, "/")
 	}
 
 	// create the client
@@ -143,73 +144,12 @@ func (p *Plugin) Exec() error {
 		client = s3.New(sess)
 	}
 
+	// If in download mode, call the downloadS3Objects method
 	if p.Download {
 		sourceDir := normalizePath(p.Source)
+		client := p.createS3Client()
 
-		log.WithFields(log.Fields{
-			"bucket": p.Bucket,
-			"dir":    sourceDir,
-		}).Info("Listing S3 directory")
-
-		list, err := client.ListObjectsV2(&s3.ListObjectsV2Input{
-			Bucket: &p.Bucket,
-			Prefix: &sourceDir,
-		})
-		if err != nil {
-			log.WithFields(log.Fields{
-				"error":  err,
-				"bucket": p.Bucket,
-				"dir":    sourceDir,
-			}).Error("Cannot list S3 directory")
-			return err
-		}
-
-		for _, item := range list.Contents {
-			log.WithFields(log.Fields{
-				"bucket": p.Bucket,
-				"key":    *item.Key,
-			}).Info("Getting S3 object")
-
-			obj, err := client.GetObject(&s3.GetObjectInput{
-				Bucket: &p.Bucket,
-				Key:    item.Key,
-			})
-			if err != nil {
-				log.WithFields(log.Fields{
-					"error":  err,
-					"bucket": p.Bucket,
-					"key":    *item.Key,
-				}).Error("Cannot get S3 object")
-				return err
-			}
-			defer obj.Body.Close()
-
-			// resolveSource takes a target directory, a target path, and a prefix to strip,
-			// and returns a resolved source path by removing the targetDir from the target
-			// and appending the stripPrefix.
-			target := resolveSource(sourceDir, *item.Key, p.StripPrefix)
-
-			f, err := os.Create(target)
-			if err != nil {
-				log.WithFields(log.Fields{
-					"error": err,
-					"file":  target,
-				}).Error("Failed to create file")
-				return err
-			}
-			defer f.Close()
-
-			_, err = io.Copy(f, obj.Body)
-			if err != nil {
-				log.WithFields(log.Fields{
-					"error": err,
-					"file":  target,
-				}).Error("Failed to write file")
-				return err
-			}
-		}
-
-		return nil
+		return p.downloadS3Objects(client, sourceDir)
 	}
 
 	// find the bucket
@@ -399,9 +339,9 @@ func resolveKey(target, srcPath, stripPrefix string) string {
 	return key
 }
 
-func resolveSource(targetDir, target, stripPrefix string) string {
-	// Remove the leading targetDir from the target path
-	path := strings.TrimPrefix(strings.TrimPrefix(target, targetDir), "/")
+func resolveSource(sourceDir, source, stripPrefix string) string {
+	// Remove the leading sourceDir from the source path
+	path := strings.TrimPrefix(strings.TrimPrefix(source, sourceDir), "/")
 
 	// Add the specified stripPrefix to the resulting path
 	return stripPrefix + path
@@ -429,6 +369,126 @@ func isDir(source string, matches []string) bool {
 }
 
 // normalizePath converts the path to a forward slash format and trims the prefix.
-func normalizePath(source string) string {
-	return strings.TrimPrefix(filepath.ToSlash(source), "/")
+func normalizePath(path string) string {
+	return strings.TrimPrefix(filepath.ToSlash(path), "/")
+}
+
+// downloadS3Object downloads a single object from S3
+func (p *Plugin) downloadS3Object(client *s3.S3, sourceDir, key, target string) error {
+	log.WithFields(log.Fields{
+		"bucket": p.Bucket,
+		"key":    key,
+	}).Info("Getting S3 object")
+
+	obj, err := client.GetObject(&s3.GetObjectInput{
+		Bucket: &p.Bucket,
+		Key:    &key,
+	})
+	if err != nil {
+		log.WithFields(log.Fields{
+			"error":  err,
+			"bucket": p.Bucket,
+			"key":    key,
+		}).Error("Cannot get S3 object")
+		return err
+	}
+	defer obj.Body.Close()
+
+	// Create the destination file path
+	destination := filepath.Join(p.Target, target)
+	log.Println("Destination: ", destination)
+
+	// Extract the directory from the destination path
+	dir := filepath.Dir(destination)
+
+	// Create the directory and any necessary parent directories
+	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
+		return errors.Wrap(err, "error creating directories")
+	}
+
+	f, err := os.Create(destination)
+	if err != nil {
+		log.WithFields(log.Fields{
+			"error": err,
+			"file":  destination,
+		}).Error("Failed to create file")
+		return err
+	}
+	defer f.Close()
+
+	_, err = io.Copy(f, obj.Body)
+	if err != nil {
+		log.WithFields(log.Fields{
+			"error": err,
+			"file":  destination,
+		}).Error("Failed to write file")
+		return err
+	}
+
+	return nil
+}
+
+// downloadS3Objects downloads all objects in the specified S3 bucket path
+func (p *Plugin) downloadS3Objects(client *s3.S3, sourceDir string) error {
+	log.WithFields(log.Fields{
+		"bucket": p.Bucket,
+		"dir":    sourceDir,
+	}).Info("Listing S3 directory")
+
+	list, err := client.ListObjectsV2(&s3.ListObjectsV2Input{
+		Bucket: &p.Bucket,
+		Prefix: &sourceDir,
+	})
+	if err != nil {
+		log.WithFields(log.Fields{
+			"error":  err,
+			"bucket": p.Bucket,
+			"dir":    sourceDir,
+		}).Error("Cannot list S3 directory")
+		return err
+	}
+
+	for _, item := range list.Contents {
+		// resolveSource takes a source directory, a source path, and a prefix to strip,
+		// and returns a resolved target path by removing the sourceDir from the source
+		// and appending the stripPrefix.
+		target := resolveSource(sourceDir, *item.Key, p.StripPrefix)
+
+		if err := p.downloadS3Object(client, sourceDir, *item.Key, target); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// 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
 }

plugin_windows_test.go

@@ -59,41 +59,6 @@ func TestResolveWinKey(t *testing.T) {
 	}
 }
 
-func TestResolveDir(t *testing.T) {
-	tests := []struct {
-		input    string
-		expected string
-	}{
-		{
-			input:    "example-string",
-			expected: "example-string",
-		},
-		{
-			input:    "/path/to/file",
-			expected: "path/to/file",
-		},
-		{
-			input:    "12345",
-			expected: "12345",
-		},
-		{
-			input:    "/root/directory",
-			expected: "root/directory",
-		},
-		{
-			input:    "no_slash",
-			expected: "no_slash",
-		},
-	}
-
-	for _, tc := range tests {
-		result := resolveDir(tc.input)
-		if result != tc.expected {
-			t.Errorf("Expected: %s, Got: %s", tc.expected, result)
-		}
-	}
-}
-
 func TestNormalizePath(t *testing.T) {
 	tests := []struct {
 		input    string