feat: improve password management (#58)

Co-authored-by: Giulio <giuliocalzolari@users.noreply.github.com>
datarobot-oss · Jan 9, 2025 · adc0532 · adc0532
1 parent 272e749
commit adc0532
Show file tree

Hide file tree

Showing 6 changed files with 132 additions and 10 deletions.
diff --git a/cmd/load.go b/cmd/load.go
@@ -30,7 +30,21 @@ Example:
 $ helm datarobot load images.tgz -r registry.example.com -u reg_username -p reg_password
 Successfully pushed image: registry.example.com/test-image1:1.0.0
 
-'''`, "'", "`", -1),
+'''
+
+Authentication can be provided in various ways, including:
+
+'''sh
+export REGISTRY_USERNAME=reg_username
+export REGISTRY_PASSWORD=reg_password
+$ helm datarobot load images.tgz -r registry.example.com
+'''
+
+'''sh
+$ echo "reg_password" | helm datarobot load images.tgz -r registry.example.com -u reg_username --password-stdin
+'''
+
+`, "'", "`", -1),
 	Args: cobra.MinimumNArgs(1), // Requires at least one argument (file path)
 	RunE: func(cmd *cobra.Command, args []string) error {
 
@@ -118,13 +132,13 @@ Successfully pushed image: registry.example.com/test-image1:1.0.0
 
 			if loadToken != "" {
 				auth = &authn.Bearer{
-					Token: loadToken,
+					Token: GetSecret(loadPasswordStdin, "REGISTRY_TOKEN", loadToken),
 				}
 
 			} else if loadUsername != "" && loadPassword != "" {
 				auth = &authn.Basic{
-					Username: loadUsername,
-					Password: loadPassword,
+					Username: GetSecret(false, "REGISTRY_USERNAME", loadUsername),
+					Password: GetSecret(loadPasswordStdin, "REGISTRY_PASSWORD", loadPassword),
 				}
 
 			} else {
@@ -143,12 +157,13 @@ Successfully pushed image: registry.example.com/test-image1:1.0.0
 }
 
 var loadReg, loadUsername, loadPassword, loadToken, loadImagePrefix, loadImageSuffix, loadImageRepo string
-var loadDryRun bool
+var loadDryRun, loadPasswordStdin bool
 
 func init() {
 	rootCmd.AddCommand(loadCmd)
 	loadCmd.Flags().StringVarP(&loadUsername, "username", "u", "", "username to auth")
 	loadCmd.Flags().StringVarP(&loadPassword, "password", "p", "", "pass to auth")
+	loadCmd.Flags().BoolVar(&loadPasswordStdin, "password-stdin", false, "Read password from stdin")
 	loadCmd.Flags().StringVarP(&loadToken, "token", "t", "", "pass to auth")
 	loadCmd.Flags().StringVarP(&loadReg, "registry", "r", "", "registry to auth")
 	loadCmd.Flags().StringVarP(&loadImagePrefix, "prefix", "", "", "append prefix on repo name")

diff --git a/cmd/sync.go b/cmd/sync.go
@@ -25,8 +25,21 @@ $ helm datarobot sync testdata/test-chart1/ -r registry.example.com -u reg_usern
 
 Pulling image: docker.io/datarobot/test-image1:1.0.0
 Pushing image: registry.example.com/datarobot/test-image1:1.0.0
+'''
 
-'''`, "'", "`", -1),
+Authentication can be provided in various ways, including:
+
+'''sh
+export REGISTRY_USERNAME=reg_username
+export REGISTRY_PASSWORD=reg_password
+$ helm datarobot sync testdata/test-chart1/ -r registry.example.com
+'''
+
+'''sh
+$ echo "reg_password" | helm datarobot sync testdata/test-chart1/ -r registry.example.com -u reg_username --password-stdin
+'''
+
+`, "'", "`", -1),
 	RunE: func(cmd *cobra.Command, args []string) error {
 		images, err := ExtractImagesFromCharts(args)
 		if err != nil {
@@ -74,13 +87,13 @@ Pushing image: registry.example.com/datarobot/test-image1:1.0.0
 
 			if syncToken != "" {
 				auth = &authn.Bearer{
-					Token: syncToken,
+					Token: GetSecret(syncPasswordStdin, "REGISTRY_TOKEN", syncToken),
 				}
 
 			} else if syncUsername != "" && syncPassword != "" {
 				auth = &authn.Basic{
-					Username: syncUsername,
-					Password: syncPassword,
+					Username: GetSecret(false, "REGISTRY_USERNAME", syncUsername),
+					Password: GetSecret(syncPasswordStdin, "REGISTRY_PASSWORD", syncPassword),
 				}
 
 			} else {
@@ -97,13 +110,14 @@ Pushing image: registry.example.com/datarobot/test-image1:1.0.0
 }
 
 var syncReg, syncUsername, syncPassword, syncToken, syncImagePrefix, syncImageSuffix, syncImageRepo, syncTransform, caCertPath, certPath, keyPath string
-var syncDryRun, skipTlsVerify bool
+var syncDryRun, skipTlsVerify, syncPasswordStdin bool
 
 func init() {
 	rootCmd.AddCommand(syncCmd)
 	syncCmd.Flags().StringVarP(&annotation, "annotation", "a", "datarobot.com/images", "annotation to lookup")
 	syncCmd.Flags().StringVarP(&syncUsername, "username", "u", "", "username to auth")
 	syncCmd.Flags().StringVarP(&syncPassword, "password", "p", "", "pass to auth")
+	syncCmd.Flags().BoolVar(&syncPasswordStdin, "password-stdin", false, "Read password from stdin")
 	syncCmd.Flags().StringVarP(&syncToken, "token", "t", "", "pass to auth")
 	syncCmd.Flags().StringVarP(&syncReg, "registry", "r", "", "registry to auth")
 	syncCmd.Flags().StringVarP(&syncImagePrefix, "prefix", "", "", "append prefix on repo name")

diff --git a/cmd/utils.go b/cmd/utils.go
@@ -1,11 +1,13 @@
 package cmd
 
 import (
+	"bufio"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"sort"
 	"strings"
 
@@ -121,3 +123,22 @@ func ExtractImagesFromManifest(manifest string) ([]string, error) {
 	sort.Strings(manifestImages)
 	return manifestImages, nil
 }
+
+func GetSecret(stdin bool, envVar string, argv string) string {
+	secret := ""
+	if stdin {
+		scanner := bufio.NewScanner(os.Stdin)
+		if scanner.Scan() {
+			secret = scanner.Text()
+		}
+	}
+	if secret == "" {
+		value, exists := os.LookupEnv(envVar)
+		if !exists {
+			secret = argv
+		} else {
+			secret = value
+		}
+	}
+	return secret
+}
diff --git a/cmd/utils_test.go b/cmd/utils_test.go
@@ -4,6 +4,8 @@ import (
 	"io/ioutil"
 	"os"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestGetTransport(t *testing.T) {
@@ -94,3 +96,44 @@ E4bmYvhnmO/hlPwDN02OSWHYm6m0yIzWXw==
 		t.Fatal("Expected InsecureSkipVerify to be true")
 	}
 }
+
+func TestGetSecret(t *testing.T) {
+	t.Run("direct", func(t *testing.T) {
+		output := GetSecret(false, "TEST_SEC", "test")
+		expectedOutput := `test`
+		assert.Equal(t, expectedOutput, output)
+	})
+	t.Run("env", func(t *testing.T) {
+		os.Setenv("TEST_SEC", "test")
+		output := GetSecret(false, "TEST_SEC", "")
+		expectedOutput := `test`
+		assert.Equal(t, expectedOutput, output)
+	})
+	t.Run("sdtin", func(t *testing.T) {
+		input := "test\n"
+		// Save the original os.Stdin
+		originalStdin := os.Stdin
+		defer func() { os.Stdin = originalStdin }()
+
+		// Create a temporary os.Stdin
+		r, w, err := os.Pipe()
+		if err != nil {
+			t.Fatalf("Failed to create pipe: %v", err)
+		}
+
+		// Write the mock input to the pipe
+		_, err = w.Write([]byte(input))
+		if err != nil {
+			t.Fatalf("Failed to write to pipe: %v", err)
+		}
+		// Close the pipe write to simulate end of input
+		w.Close()
+
+		// Redirect os.Stdin to read from the pipe
+		os.Stdin = r
+		output := GetSecret(true, "", "")
+		expectedOutput := `test`
+		assert.Equal(t, expectedOutput, output)
+	})
+
+}
diff --git a/docs/helm-datarobot_load.md b/docs/helm-datarobot_load.md
@@ -15,6 +15,20 @@ Successfully pushed image: registry.example.com/test-image1:1.0.0
 
 ```
 
+Authentication can be provided in various ways, including:
+
+```sh
+export REGISTRY_USERNAME=reg_username
+export REGISTRY_PASSWORD=reg_password
+$ helm datarobot load images.tgz -r registry.example.com
+```
+
+```sh
+$ echo "reg_password" | helm datarobot load images.tgz -r registry.example.com -u reg_username --password-stdin
+```
+
+
+
 ```
 helm-datarobot load [flags]
 ```
@@ -29,6 +43,7 @@ helm-datarobot load [flags]
   -i, --insecure          Skip server certificate verification
   -K, --key string        Path to the client key
   -p, --password string   pass to auth
+      --password-stdin    Read password from stdin
       --prefix string     append prefix on repo name
   -r, --registry string   registry to auth
       --repo string       rewrite the target repository name

diff --git a/docs/helm-datarobot_sync.md b/docs/helm-datarobot_sync.md
@@ -14,9 +14,22 @@ $ helm datarobot sync testdata/test-chart1/ -r registry.example.com -u reg_usern
 
 Pulling image: docker.io/datarobot/test-image1:1.0.0
 Pushing image: registry.example.com/datarobot/test-image1:1.0.0
+```
+
+Authentication can be provided in various ways, including:
+
+```sh
+export REGISTRY_USERNAME=reg_username
+export REGISTRY_PASSWORD=reg_password
+$ helm datarobot sync testdata/test-chart1/ -r registry.example.com
+```
 
+```sh
+$ echo "reg_password" | helm datarobot sync testdata/test-chart1/ -r registry.example.com -u reg_username --password-stdin
 ```
 
+
+
 ```
 helm-datarobot sync [flags]
 ```
@@ -32,6 +45,7 @@ helm-datarobot sync [flags]
   -i, --insecure            Skip server certificate verification
   -K, --key string          Path to the client key
   -p, --password string     pass to auth
+      --password-stdin      Read password from stdin
       --prefix string       append prefix on repo name
   -r, --registry string     registry to auth
       --repo string         rewrite the target repository name