From e2561f9073cee846f28b6f024d937cbb0bf221fd Mon Sep 17 00:00:00 2001
From: Steve Kriss <steve@heptio.com>
Date: Thu, 7 Jun 2018 09:52:55 -0700
Subject: [PATCH] ark restic init-repository: use key data if provided

Signed-off-by: Steve Kriss <steve@heptio.com>
---
 pkg/cmd/cli/restic/init_repository.go      | 20 ++++---
 pkg/cmd/cli/restic/init_repository_test.go | 67 ++++++++++++++++++++++
 2 files changed, 80 insertions(+), 7 deletions(-)
 create mode 100644 pkg/cmd/cli/restic/init_repository_test.go

diff --git a/pkg/cmd/cli/restic/init_repository.go b/pkg/cmd/cli/restic/init_repository.go
index cfeb0dff2..51b1b8b29 100644
--- a/pkg/cmd/cli/restic/init_repository.go
+++ b/pkg/cmd/cli/restic/init_repository.go
@@ -65,6 +65,11 @@ func NewInitRepositoryOptions() *InitRepositoryOptions {
 	}
 }
 
+var (
+	errKeyFileAndKeyDataProvided = errors.Errorf("only one of --key-file and --key-data may be specified")
+	errKeySizeTooSmall           = errors.Errorf("--key-size must be at least 1")
+)
+
 func (o *InitRepositoryOptions) BindFlags(flags *pflag.FlagSet) {
 	flags.StringVar(&o.KeyFile, "key-file", o.KeyFile, "Path to file containing the encryption key for the restic repository. Optional; if unset, Ark will generate a random key for you.")
 	flags.StringVar(&o.KeyData, "key-data", o.KeyData, "Encryption key for the restic repository. Optional; if unset, Ark will generate a random key for you.")
@@ -73,27 +78,28 @@ func (o *InitRepositoryOptions) BindFlags(flags *pflag.FlagSet) {
 
 func (o *InitRepositoryOptions) Complete(f client.Factory) error {
 	if o.KeyFile != "" && o.KeyData != "" {
-		return errors.Errorf("only one of --key-file and --key-data may be specified")
+		return errKeyFileAndKeyDataProvided
 	}
 
 	if o.KeyFile == "" && o.KeyData == "" && o.KeySize < 1 {
-		return errors.Errorf("--key-size must be at least 1")
+		return errKeySizeTooSmall
 	}
 
 	o.Namespace = f.Namespace()
 
-	if o.KeyFile != "" {
+	switch {
+	case o.KeyFile != "":
 		data, err := ioutil.ReadFile(o.KeyFile)
 		if err != nil {
 			return err
 		}
 		o.keyBytes = data
-	}
-
-	if len(o.KeyData) == 0 {
+	case o.KeyData != "":
+		o.keyBytes = []byte(o.KeyData)
+	case o.KeySize > 0:
 		o.keyBytes = make([]byte, o.KeySize)
 		// rand.Reader always returns a nil error
-		_, _ = rand.Read(o.keyBytes)
+		rand.Read(o.keyBytes)
 	}
 
 	return nil
diff --git a/pkg/cmd/cli/restic/init_repository_test.go b/pkg/cmd/cli/restic/init_repository_test.go
new file mode 100644
index 000000000..1b38eaefb
--- /dev/null
+++ b/pkg/cmd/cli/restic/init_repository_test.go
@@ -0,0 +1,67 @@
+package restic
+
+import (
+	"testing"
+
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/assert"
+
+	"k8s.io/client-go/kubernetes"
+
+	"github.com/heptio/ark/pkg/client"
+	clientset "github.com/heptio/ark/pkg/generated/clientset/versioned"
+)
+
+type fakeFactory struct{}
+
+var _ client.Factory = &fakeFactory{}
+
+func (f *fakeFactory) BindFlags(flags *pflag.FlagSet) {
+	panic("not implemented")
+}
+
+func (f *fakeFactory) Client() (clientset.Interface, error) {
+	panic("not implemented")
+}
+
+func (f *fakeFactory) KubeClient() (kubernetes.Interface, error) {
+	panic("not implemented")
+}
+
+func (f *fakeFactory) Namespace() string {
+	return ""
+}
+
+func TestComplete(t *testing.T) {
+	// no key options provided should error
+	o := &InitRepositoryOptions{}
+	err := o.Complete(&fakeFactory{})
+	assert.EqualError(t, err, errKeySizeTooSmall.Error())
+
+	// both KeyFile and KeyData provided should error
+	o = &InitRepositoryOptions{
+		KeyFile: "/foo",
+		KeyData: "bar",
+	}
+	err = o.Complete(&fakeFactory{})
+	assert.EqualError(t, err, errKeyFileAndKeyDataProvided.Error())
+
+	// TODO test that if KeyFile is provided, the data is used
+	// (should move the pkg/restore/FileSystem interface to
+	// a common location and use it here so ioutil.ReadFile
+	// can be easily mocked)
+
+	// if KeyData is provided, it's used
+	o = &InitRepositoryOptions{
+		KeyData: "bar",
+	}
+	assert.NoError(t, o.Complete(&fakeFactory{}))
+	assert.Equal(t, []byte(o.KeyData), o.keyBytes)
+
+	// if KeySize is provided, a random key is generated
+	o = &InitRepositoryOptions{
+		KeySize: 10,
+	}
+	assert.NoError(t, o.Complete(&fakeFactory{}))
+	assert.Equal(t, o.KeySize, len(o.keyBytes))
+}