From 702279607073958ea4bc4fd8b6552ac2ed727f49 Mon Sep 17 00:00:00 2001
From: Andrew Lavery <laverya@umich.edu>
Date: Thu, 3 Jan 2019 17:04:34 -0800
Subject: [PATCH] allow specifying the charset to use with the RandomString
 template function

---
 pkg/templates/random.go      |  9 ++++-
 pkg/templates/random_test.go | 77 ++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/pkg/templates/random.go b/pkg/templates/random.go
index 92af28c1c..7f082cbba 100644
--- a/pkg/templates/random.go
+++ b/pkg/templates/random.go
@@ -16,10 +16,15 @@ const (
 )
 
 // stolen from https://github.com/replicatedhq/replicated/blob/8ce3ed40436e38b8089387d103623dbe09bbf1c0/pkg/commands/random.go#L22
-func (ctx *StaticCtx) RandomString(length uint64) string {
+func (ctx *StaticCtx) RandomString(length uint64, providedCharset ...string) string {
 	debug := log.With(level.Debug(ctx.Logger), "func", "random")
 	debug.Log("event", "start")
-	regExp, err := syntax.Parse(DefaultCharset, syntax.Perl)
+
+	charset := DefaultCharset
+	if len(providedCharset) >= 1 {
+		charset = providedCharset[0]
+	}
+	regExp, err := syntax.Parse(charset, syntax.Perl)
 	if err != nil {
 		debug.Log("event", "regexParse.fail", "err", err)
 		return ""
diff --git a/pkg/templates/random_test.go b/pkg/templates/random_test.go
index 00e874717..4e51c98c6 100644
--- a/pkg/templates/random_test.go
+++ b/pkg/templates/random_test.go
@@ -1,10 +1,15 @@
 package templates
 
 import (
+	"fmt"
+	"regexp"
 	"testing"
+	"unicode/utf8"
 
 	"github.com/replicatedhq/ship/pkg/testing/logger"
+	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestGenerateRandomString(t *testing.T) {
@@ -14,3 +19,75 @@ func TestGenerateRandomString(t *testing.T) {
 	assert.Regexp(t, DefaultCharset, str)
 
 }
+
+func TestGenerateRandomStringTemplates(t *testing.T) {
+	tests := []struct {
+		name           string
+		length         int
+		outputRegex    string
+		templateString string
+	}{
+		{
+			name:           "DefaultCharset",
+			length:         100,
+			outputRegex:    DefaultCharset,
+			templateString: `{{repl RandomString 100}}`,
+		},
+		{
+			name:           "ExplicitDefaultCharset",
+			length:         100,
+			outputRegex:    DefaultCharset,
+			templateString: fmt.Sprintf("{{repl RandomString 100 %q}}", DefaultCharset),
+		},
+		{
+			name:           "lowercase",
+			length:         100,
+			outputRegex:    "[a-z]",
+			templateString: `{{repl RandomString 100 "[a-z]"}}`,
+		},
+		{
+			name:           "very restricted charset aB",
+			length:         100,
+			outputRegex:    "[aB]",
+			templateString: `{{repl RandomString 100 "[aB]"}}`,
+		},
+		{
+			name:           "not lowercase",
+			length:         1000,
+			outputRegex:    "[^a-z]",
+			templateString: `{{repl RandomString 1000 "[^a-z]"}}`,
+		},
+		{
+			name:           "base64",
+			length:         1000,
+			outputRegex:    "[A-Za-z0-9+/]",
+			templateString: `{{repl RandomString 1000 "[A-Za-z0-9+/]"}}`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := require.New(t)
+
+			builderBuilder := &BuilderBuilder{
+				Logger: &logger.TestLogger{T: t},
+				Viper:  viper.New(),
+			}
+
+			builder := builderBuilder.NewBuilder(
+				builderBuilder.NewStaticContext(),
+			)
+
+			outputString, err := builder.String(tt.templateString)
+			req.NoError(err)
+
+			req.Equal(utf8.RuneCountInString(outputString), tt.length, "%q should have length %d", outputString, tt.length)
+
+			tt.outputRegex = fmt.Sprintf("^%s+$", tt.outputRegex)
+
+			matcher := regexp.MustCompile(tt.outputRegex)
+
+			req.True(matcher.MatchString(outputString), "String %q must match %q", outputString, tt.outputRegex)
+		})
+	}
+}