Repository: sa7mon/S3Scanner (language: python)

Issue:
De-dupe bucket list when ingesting bucket file
Implement `set` like Python?


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index df65d06..2ac8473 100644
--- a/README.md
+++ b/README.md
@@ -137,6 +137,8 @@ assets
 image-uploads
 ```
 
+Bucket names listed multiple times will only be scanned once.
+
 *`-mq`*
 -------
 
diff --git a/bucket/bucket.go b/bucket/bucket.go
index 001b77a..a529e86 100644
--- a/bucket/bucket.go
+++ b/bucket/bucket.go
@@ -3,15 +3,17 @@ package bucket
 import (
 	"bufio"
 	"fmt"
-	"github.com/aws/aws-sdk-go-v2/service/s3"
-	"github.com/aws/aws-sdk-go-v2/service/s3/types"
-	"github.com/sa7mon/s3scanner/groups"
-	log "github.com/sirupsen/logrus"
+	"io"
 	"os"
 	"regexp"
 	"strings"
 	"time"
 	"unicode"
+
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/aws/aws-sdk-go-v2/service/s3/types"
+	"github.com/sa7mon/s3scanner/groups"
+	log "github.com/sirupsen/logrus"
 )
 
 var BucketExists = uint8(1)
@@ -148,28 +150,40 @@ func (bucket *Bucket) Permissions() map[*types.Grantee]map[string]uint8 {
 	}
 }
 
-func ReadFromFile(bucketFile string, bucketChan chan Bucket) error {
-	file, err := os.Open(bucketFile)
-	if err != nil {
-		return err
-	}
-	defer file.Close()
-
-	fileScanner := bufio.NewScanner(file)
-	for fileScanner.Scan() {
-		bucketName := strings.TrimSpace(fileScanner.Text())
+func FromReader(r io.Reader, bucketChan chan Bucket) error {
+	scanner := bufio.NewScanner(r)
+	bucketsSeen := make(map[string]struct{})
+	for scanner.Scan() {
+		bucketName := strings.TrimSpace(scanner.Text())
 		if !IsValidS3BucketName(bucketName) {
 			log.Info(fmt.Sprintf("invalid   | %s", bucketName))
-		} else {
-			bucketChan <- NewBucket(strings.ToLower(bucketName))
+			continue
+		}
+		bucketName = strings.ToLower(bucketName)
+		if _, seen := bucketsSeen[bucketName]; seen {
+			continue
 		}
+		bucketsSeen[bucketName] = struct{}{}
+		bucketChan <- NewBucket(bucketName)
 	}
 
-	if ferr := fileScanner.Err(); ferr != nil {
+	if ferr := scanner.Err(); ferr != nil {
 		return ferr
 	}
+	return nil
+}
+
+func ReadFromFile(bucketFile string, bucketChan chan Bucket) error {
+	file, err := os.Open(bucketFile)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
 
-	return err
+	if err := FromReader(file, bucketChan); err != nil {
+		return err
+	}
+	return nil
 }
 
 // ParseAclOutputv2 TODO: probably move this to providers.go


Test patch:
diff --git a/bucket/bucket_test.go b/bucket/bucket_test.go
index f8e6386..6d8c3c8 100644
--- a/bucket/bucket_test.go
+++ b/bucket/bucket_test.go
@@ -3,15 +3,17 @@ package bucket
 import (
 	"context"
 	"fmt"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"testing"
+	"time"
+
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/aws/aws-sdk-go-v2/service/s3/types"
 	"github.com/sa7mon/s3scanner/groups"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/sync/errgroup"
-	"path/filepath"
-	"runtime"
-	"testing"
-	"time"
 )
 
 func TestIsValidS3BucketName_Good(t *testing.T) {
@@ -262,6 +264,41 @@ func TestBucket_ParseAclOutputv2(t *testing.T) {
 	}
 }
 
+func TestFromReader(t *testing.T) {
+	t.Parallel()
+
+	reader := strings.NewReader(`test
+bar
+bucket
+bar
+test
+foo
+bucket
+foo
+bar`)
+
+	testChan := make(chan Bucket)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	g, _ := errgroup.WithContext(ctx)
+	defer cancel()
+
+	g.Go(func() error {
+		err := FromReader(reader, testChan)
+		close(testChan)
+		return err
+	})
+
+	i := 0
+	for range testChan {
+		i++
+	}
+	assert.Equal(t, 4, i)
+
+	if err := g.Wait(); err != nil {
+		t.Error(err)
+	}
+}
+
 func TestReadFromFile(t *testing.T) {
 	t.Parallel()