From 269bb2691aa9d31652700caed62fa080793efdb3 Mon Sep 17 00:00:00 2001
From: Cthulhu <cthulhu.den@gmail.com>
Date: Wed, 28 Sep 2022 02:02:32 +0300
Subject: [PATCH] Add ability to download the files back from S3

---
 go.mod    |  1 +
 go.sum    |  2 ++
 main.go   |  6 +++++
 plugin.go | 80 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 89 insertions(+)

diff --git a/go.mod b/go.mod
index f2c8b37..30b977f 100644
--- a/go.mod
+++ b/go.mod
@@ -15,5 +15,6 @@ require (
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd // indirect
+	golang.org/x/sync v0.0.0-20220923202941-7f9b1623fab7 // indirect
 	golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect
 )
diff --git a/go.sum b/go.sum
index 52fc6c9..d4ce394 100644
--- a/go.sum
+++ b/go.sum
@@ -48,6 +48,8 @@ golang.org/x/net v0.0.0-20210614182718-04defd469f4e/go.mod h1:9nx3DQGgdP8bBQD5qx
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2 h1:CIJ76btIcR3eFI5EgSo6k1qKw9KJexJuRLI9G7Hp5wE=
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
+golang.org/x/sync v0.0.0-20220923202941-7f9b1623fab7 h1:ZrnxWX62AgTKOSagEqxvb3ffipvEDX2pl7E1TdqLqIc=
+golang.org/x/sync v0.0.0-20220923202941-7f9b1623fab7/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
diff --git a/main.go b/main.go
index cc2c333..ec61a71 100644
--- a/main.go
+++ b/main.go
@@ -95,6 +95,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 {
 		CacheControl:          c.Generic("cache-control").(*StringMapFlag).Get(),
 		StorageClass:          c.String("storage-class"),
 		PathStyle:             c.Bool("path-style"),
+		Download:              c.Bool("download"),
 		DryRun:                c.Bool("dry-run"),
 	}
 
diff --git a/plugin.go b/plugin.go
index a812d49..d1141c3 100644
--- a/plugin.go
+++ b/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.
@@ -88,6 +90,8 @@ type Plugin struct {
 	//
 	// Should be true for minio and false for AWS.
 	PathStyle bool
+	// if true, plugin is set to download mode, which means `target` from the bucket will be downloaded
+	Download bool
 	// Dry run without uploading/
 	DryRun bool
 }
@@ -134,6 +138,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,
@@ -318,3 +393,8 @@ 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
+}