Download function added

Author: shiv-am0

go.mod

@@ -14,5 +14,6 @@ require (
 	github.com/cpuguy83/go-md2man/v2 v2.0.1 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	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

@@ -40,6 +40,8 @@ golang.org/x/net v0.1.0 h1:hZ/3BUoy5aId7sCpA/Tc5lt8DkFgdVS2onTpJsZ/fl0=
 golang.org/x/net v0.1.0/go.mod h1:Cx3nUiGt4eDBEyega/BKRp+/AlGL8hYe7U9odMt2Cco=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.6.0 h1:5BMeUDZ7vkXGfEr1x9B4bRcTH4lpkTkpdh0T/J+qjbQ=
+golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

main.go

@@ -94,6 +94,11 @@ func main() {
 			Usage:  "server-side encryption algorithm, defaults to none",
 			EnvVar: "PLUGIN_ENCRYPTION",
 		},
+		cli.BoolFlag{
+			Name:   "download",
+			Usage:  "switch to download mode, which will fetch `target`'s files from s3 bucket and place them according to `strip-prefix`",
+			EnvVar: "PLUGIN_DOWNLOAD",
+		},
 		cli.BoolFlag{
 			Name:   "dry-run",
 			Usage:  "dry run for debug purposes",
@@ -164,6 +169,7 @@ func run(c *cli.Context) error {
 		Exclude:               c.StringSlice("exclude"),
 		Encryption:            c.String("encryption"),
 		ContentType:           c.Generic("content-type").(*StringMapFlag).Get(),
+		Download:              c.Bool("download"),
 		ContentEncoding:       c.Generic("content-encoding").(*StringMapFlag).Get(),
 		CacheControl:          c.Generic("cache-control").(*StringMapFlag).Get(),
 		StorageClass:          c.String("storage-class"),

plugin.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"io"
 	"mime"
 	"os"
 	"path/filepath"
@@ -16,6 +17,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/sts"
 	"github.com/mattn/go-zglob"
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
 )
 
 // Plugin defines the S3 plugin parameters.
@@ -44,6 +46,9 @@ type Plugin struct {
 	// sa-east-1
 	Region string
 
+	// if true, plugin is set to download mode, which means `target` from the bucket will be downloaded
+	Download bool
+
 	// Indicates the files ACL, which should be one
 	// of the following:
 	//     private
@@ -135,6 +140,77 @@ func (p *Plugin) Exec() error {
 		client = s3.New(sess)
 	}
 
+	if p.Download {
+		targetDir := strings.TrimPrefix(filepath.ToSlash(p.Target), "/")
+		log.WithFields(log.Fields{
+			"bucket": p.Bucket,
+			"dir":    targetDir,
+		}).Info("Listing S3 directory")
+
+		list, err := client.ListObjectsV2(&s3.ListObjectsV2Input{
+			Bucket: &p.Bucket,
+			Prefix: &targetDir,
+		})
+		if err != nil {
+			log.WithFields(log.Fields{
+				"error":  err,
+				"bucket": p.Bucket,
+				"dir":    targetDir,
+			}).Error("Cannot list S3 directory")
+			return err
+		}
+
+		g := errgroup.Group{}
+
+		for _, item := range list.Contents {
+			log.WithFields(log.Fields{
+				"bucket": p.Bucket,
+				"key":    *item.Key,
+			}).Info("Getting S3 object")
+
+			item := item
+			g.Go(func() error {
+				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
+				}
+
+				source := resolveSource(targetDir, *item.Key, p.StripPrefix)
+
+				f, err := os.Create(source)
+				if err != nil {
+					log.WithFields(log.Fields{
+						"error": err,
+						"file":  source,
+					}).Error("Problem opening file for writing")
+					return err
+				}
+				defer f.Close()
+
+				_, err = io.Copy(f, obj.Body)
+				if err != nil {
+					log.WithFields(log.Fields{
+						"error": err,
+						"file":  source,
+					}).Error("Failed to write file")
+					return err
+				}
+
+				return nil
+			})
+		}
+
+		return g.Wait()
+	}
+
 	// find the bucket
 	log.WithFields(log.Fields{
 		"region":   p.Region,
@@ -322,6 +398,11 @@ func resolveKey(target, srcPath, stripPrefix string) string {
 	return key
 }
 
+func resolveSource(targetDir, target, stripPrefix string) string {
+	path := strings.TrimPrefix(strings.TrimPrefix(target, targetDir), "/")
+	return stripPrefix + path
+}
+
 // checks if the source path is a dir
 func isDir(source string, matches []string) bool {
 	stat, err := os.Stat(source)