Repository: datreeio/datree
Pull request #472 (2022-03-08 15:34:58)

feat: add create token and client id set on local config


Patch:
diff --git a/bl/errorReporter/errorReporter.go b/bl/errorReporter/errorReporter.go
index 0ec7918..c0b7131 100644
--- a/bl/errorReporter/errorReporter.go
+++ b/bl/errorReporter/errorReporter.go
@@ -13,17 +13,17 @@ import (
 )
 
 func ReportCliPanicError(panicErr interface{}) {
-	reporter := NewErrorReporter(cliClient.NewCliClient(deploymentConfig.URL), localConfig.NewLocalConfig(), printer.CreateNewPrinter())
+	reporter := NewErrorReporter(cliClient.NewCliClient(deploymentConfig.URL), localConfig.NewLocalConfigClient(&cliClient.CliClient{}), printer.CreateNewPrinter())
 	reporter.ReportCliError(panicErr, "/report-cli-panic-error")
 }
 
 func ReportCliUnexpectedError(unexpectedError error) {
-	reporter := NewErrorReporter(cliClient.NewCliClient(deploymentConfig.URL), localConfig.NewLocalConfig(), printer.CreateNewPrinter())
+	reporter := NewErrorReporter(cliClient.NewCliClient(deploymentConfig.URL), localConfig.NewLocalConfigClient(&cliClient.CliClient{}), printer.CreateNewPrinter())
 	reporter.ReportCliError(unexpectedError, "/report-cli-unexpected-error")
 }
 
 type LocalConfig interface {
-	GetLocalConfiguration() (*localConfig.ConfigContent, error)
+	GetLocalConfiguration() (*localConfig.LocalConfig, error)
 }
 
 type CliClient interface {
@@ -75,7 +75,7 @@ func (reporter *ErrorReporter) getCliId() (cliId string) {
 	if err != nil {
 		return "unknown"
 	} else {
-		return config.CliId
+		return config.Token
 	}
 }
 
diff --git a/cmd/config/main.go b/cmd/config/main.go
index c3e5e87..a49f031 100644
--- a/cmd/config/main.go
+++ b/cmd/config/main.go
@@ -16,7 +16,7 @@ type Printer interface {
 }
 
 type LocalConfig interface {
-	GetLocalConfiguration() (*localConfig.ConfigContent, error)
+	GetLocalConfiguration() (*localConfig.LocalConfig, error)
 	Set(key string, value string) error
 }
 
diff --git a/cmd/config/set.go b/cmd/config/set.go
index fc9967f..959f4e2 100644
--- a/cmd/config/set.go
+++ b/cmd/config/set.go
@@ -46,11 +46,6 @@ func NewSetCommand(ctx *ConfigCommandContext) *cobra.Command {
 }
 
 func set(ctx *ConfigCommandContext, key string, value string) error {
-	_, err := ctx.LocalConfig.GetLocalConfiguration()
-	if err != nil {
-		return err
-	}
-
-	err = ctx.LocalConfig.Set(key, value)
+	err := ctx.LocalConfig.Set(key, value)
 	return err
 }
diff --git a/cmd/kustomize/main.go b/cmd/kustomize/main.go
index cbf05e5..35b339e 100644
--- a/cmd/kustomize/main.go
+++ b/cmd/kustomize/main.go
@@ -65,7 +65,7 @@ func New(testCtx *test.TestCommandContext, kustomizeCtx *KustomizeContext) *cobr
 				return err
 			}
 
-			evaluationPrerunData, err := testCtx.CliClient.RequestEvaluationPrerunData(localConfigContent.CliId)
+			evaluationPrerunData, err := testCtx.CliClient.RequestEvaluationPrerunData(localConfigContent.Token)
 
 			if err != nil {
 				return err
diff --git a/cmd/publish/main.go b/cmd/publish/main.go
index 13f2ba1..90bc46e 100644
--- a/cmd/publish/main.go
+++ b/cmd/publish/main.go
@@ -20,7 +20,7 @@ type Printer interface {
 }
 
 type LocalConfig interface {
-	GetLocalConfiguration() (*localConfig.ConfigContent, error)
+	GetLocalConfiguration() (*localConfig.LocalConfig, error)
 }
 
 type CliClient interface {
@@ -36,7 +36,7 @@ type PublishCommandContext struct {
 }
 
 func New(ctx *PublishCommandContext) *cobra.Command {
-	var localConfigContent *localConfig.ConfigContent
+	var localConfigContent *localConfig.LocalConfig
 
 	publishCommand := &cobra.Command{
 		Use:   "publish <fileName>",
@@ -102,11 +102,11 @@ type MessagesContext struct {
 	CliClient   *cliClient.CliClient
 }
 
-func publish(ctx *PublishCommandContext, path string, localConfigContent *localConfig.ConfigContent) (*cliClient.PublishFailedResponse, error) {
+func publish(ctx *PublishCommandContext, path string, localConfigContent *localConfig.LocalConfig) (*cliClient.PublishFailedResponse, error) {
 	policiesConfiguration, err := files.ExtractYamlFileToUnknownStruct(path)
 	if err != nil {
 		return nil, err
 	}
 
-	return ctx.PublishCliClient.PublishPolicies(policiesConfiguration, localConfigContent.CliId)
+	return ctx.PublishCliClient.PublishPolicies(policiesConfiguration, localConfigContent.Token)
 }
diff --git a/cmd/root.go b/cmd/root.go
index 96c6cb4..2d84978 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -88,7 +88,7 @@ func Execute() error {
 }
 
 type context struct {
-	LocalConfig         *localConfig.LocalConfig
+	LocalConfig         *localConfig.LocalConfigClient
 	Evaluator           *evaluation.Evaluator
 	CliClient           *cliClient.CliClient
 	Messager            *messager.Messager
@@ -104,8 +104,8 @@ type app struct {
 }
 
 func startup() *app {
-	localConfig := localConfig.NewLocalConfig()
 	cliClient := cliClient.NewCliClient(deploymentConfig.URL)
+	localConfig := localConfig.NewLocalConfigClient(cliClient)
 	printer := printer.CreateNewPrinter()
 
 	return &app{
diff --git a/pkg/cliClient/tokens.go b/pkg/cliClient/tokens.go
new file mode 100644
index 0000000..97221c8
--- /dev/null
+++ b/pkg/cliClient/tokens.go
@@ -0,0 +1,28 @@
+package cliClient
+
+import (
+	"encoding/json"
+	"net/http"
+)
+
+type CreateTokenResponse struct {
+	Token string `json:"token"`
+}
+
+func (c *CliClient) CreateToken() (*CreateTokenResponse, error) {
+	headers := map[string]string{}
+	res, err := c.httpClient.Request(http.MethodPost, "/cli/tokens/", nil, headers)
+
+	if err != nil {
+		return nil, err
+	}
+
+	createTokenResponse := &CreateTokenResponse{}
+	err = json.Unmarshal(res.Body, &createTokenResponse)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return createTokenResponse, nil
+}
diff --git a/pkg/localConfig/localConfig.go b/pkg/localConfig/localConfig.go
index 7e20fde..2334948 100644
--- a/pkg/localConfig/localConfig.go
+++ b/pkg/localConfig/localConfig.go
@@ -5,69 +5,127 @@ import (
 	"os/user"
 	"path/filepath"
 
+	"github.com/datreeio/datree/pkg/cliClient"
 	"github.com/lithammer/shortuuid"
 	"github.com/spf13/viper"
 )
 
-type ConfigContent struct {
-	CliId         string
+type LocalConfig struct {
+	Token         string
+	ClientId      string
 	SchemaVersion string
 }
 
-type LocalConfig struct {
+type TokenClient interface {
+	CreateToken() (*cliClient.CreateTokenResponse, error)
 }
 
-func NewLocalConfig() *LocalConfig {
-	return &LocalConfig{}
+type LocalConfigClient struct {
+	tokenClient TokenClient
+}
+
+func NewLocalConfigClient(t TokenClient) *LocalConfigClient {
+	return &LocalConfigClient{
+		tokenClient: t,
+	}
 }
 
-func (lc *LocalConfig) GetLocalConfiguration() (*ConfigContent, error) {
+const (
+	clientIdKey      = "client_id"
+	tokenKey         = "token"
+	schemaVersionKey = "schema_version"
+)
+
+func (lc *LocalConfigClient) GetLocalConfiguration() (*LocalConfig, error) {
 	viper.SetEnvPrefix("datree")
 	viper.AutomaticEnv()
-	token := viper.GetString("token")
-	schemaVersion := viper.GetString("schema_version")
+
+	initConfigFileErr := InitLocalConfigFile()
+	if initConfigFileErr != nil {
+		return nil, initConfigFileErr
+	}
+
+	token := viper.GetString(tokenKey)
+	clientId := viper.GetString(clientIdKey)
+	schemaVersion := viper.GetString(schemaVersionKey)
 
 	if token == "" {
-		configHome, configName, configType, err := setViperConfig()
+		createTokenResponse, err := lc.tokenClient.CreateToken()
 		if err != nil {
 			return nil, err
 		}
-
-		// workaround for creating config file when not exist
-		// open issue in viper: https://github.com/spf13/viper/issues/430
-		// should be fixed in pr https://github.com/spf13/viper/pull/936
-		configPath := filepath.Join(configHome, configName+"."+configType)
-		_, err = os.Stat(configPath)
-
-		if err != nil {
-			os.Mkdir(configHome, os.ModePerm)
-			os.Create(configPath)
-			viper.SetDefault("token", shortuuid.New())
-			viper.WriteConfig()
+		token = createTokenResponse.Token
+		viper.SetDefault(tokenKey, token)
+		writeTokenErr := viper.WriteConfig()
+		if writeTokenErr != nil {
+			return nil, writeTokenErr
 		}
+		readTokenErr := viper.ReadInConfig()
+		if readTokenErr != nil {
+			return nil, readTokenErr
+		}
+		token = viper.GetString(tokenKey)
+	}
 
-		viper.ReadInConfig()
-		token = viper.GetString("token")
-
-		if token == "" {
-			viper.SetDefault("token", shortuuid.New())
-			viper.WriteConfig()
-			viper.ReadInConfig()
-			token = viper.GetString("token")
+	if clientId == "" {
+		viper.SetDefault(clientIdKey, shortuuid.New())
+		writeClientIdErr := viper.WriteConfig()
+		if writeClientIdErr != nil {
+			return nil, writeClientIdErr
+		}
+		readClientIdErr := viper.ReadInConfig()
+		if readClientIdErr != nil {
+			return nil, readClientIdErr
 		}
+		clientId = viper.GetString(clientIdKey)
+	}
+	return &LocalConfig{Token: token, ClientId: clientId, SchemaVersion: schemaVersion}, nil
+}
+
+func (lc *LocalConfigClient) Set(key string, value string) error {
+	initConfigFileErr := InitLocalConfigFile()
+	if initConfigFileErr != nil {
+		return initConfigFileErr
 	}
 
-	return &ConfigContent{CliId: token, SchemaVersion: schemaVersion}, nil
+	viper.Set(key, value)
+	writeClientIdErr := viper.WriteConfig()
+	if writeClientIdErr != nil {
+		return writeClientIdErr
+	}
+	return nil
 }
 
-func (lc *LocalConfig) Set(key string, value string) error {
-	_, _, _, err := setViperConfig()
+func InitLocalConfigFile() error {
+	configHome, configName, configType, err := setViperConfig()
 	if err != nil {
 		return err
 	}
+	// workaround for creating config file when not exist
+	// open issue in viper: https://github.com/spf13/viper/issues/430
+	// should be fixed in pr https://github.com/spf13/viper/pull/936
+	configPath := filepath.Join(configHome, configName+"."+configType)
 
-	viper.Set(key, value)
-	viper.WriteConfig()
+	_, err = os.Stat(configHome)
+	if err != nil {
+		osMkdirErr := os.Mkdir(configHome, os.ModePerm)
+		if osMkdirErr != nil {
+			return osMkdirErr
+		}
+	}
+
+	_, err = os.Stat(configPath)
+	if err != nil {
+		_, osCreateErr := os.Create(configPath)
+		if osCreateErr != nil {
+			return osCreateErr
+		}
+	}
+
+	readLocalFileErr := viper.ReadInConfig()
+	if readLocalFileErr != nil {
+		return readLocalFileErr
+	}
 	return nil
 }
 


Test patch:
diff --git a/bl/errorReporter/errorReporter_test.go b/bl/errorReporter/errorReporter_test.go
index d825eda..6543f0e 100644
--- a/bl/errorReporter/errorReporter_test.go
+++ b/bl/errorReporter/errorReporter_test.go
@@ -27,10 +27,10 @@ type mockConfig struct {
 	mock.Mock
 }
 
-func (lc *mockConfig) GetLocalConfiguration() (*localConfig.ConfigContent, error) {
+func (lc *mockConfig) GetLocalConfiguration() (*localConfig.LocalConfig, error) {
 	lc.Called()
-	return &localConfig.ConfigContent{
-		CliId:         "2qRg9jzJGcA73ftqEcXuBp",
+	return &localConfig.LocalConfig{
+		Token:         "2qRg9jzJGcA73ftqEcXuBp",
 		SchemaVersion: "111",
 	}, nil
 }
diff --git a/cmd/config/set_test.go b/cmd/config/set_test.go
index 7b331af..d161d8f 100644
--- a/cmd/config/set_test.go
+++ b/cmd/config/set_test.go
@@ -55,9 +55,9 @@ type LocalConfigMock struct {
 	mock.Mock
 }
 
-func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.ConfigContent, error) {
+func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.LocalConfig, error) {
 	lc.Called()
-	return &localConfig.ConfigContent{CliId: "previousToken"}, nil
+	return &localConfig.LocalConfig{Token: "previousToken"}, nil
 }
 
 func (lc *LocalConfigMock) Set(key string, value string) error {
@@ -76,7 +76,7 @@ func TestSetCommand(t *testing.T) {
 	printerMock.On("PrintEvaluationSummary", mock.Anything, mock.Anything)
 
 	localConfigMock := &LocalConfigMock{}
-	localConfigMock.On("GetLocalConfiguration").Return(&localConfig.ConfigContent{CliId: "previousToken"}, nil)
+	localConfigMock.On("GetLocalConfiguration").Return(&localConfig.LocalConfig{Token: "previousToken"}, nil)
 	localConfigMock.On("Set", mock.Anything, mock.Anything)
 
 	ctx := &ConfigCommandContext{
diff --git a/cmd/kustomize/main_test.go b/cmd/kustomize/main_test.go
index b2ea691..9999458 100644
--- a/cmd/kustomize/main_test.go
+++ b/cmd/kustomize/main_test.go
@@ -171,7 +171,7 @@ type LocalConfigMock struct {
 	mock.Mock
 }
 
-func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.ConfigContent, error) {
+func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.LocalConfig, error) {
 	lc.Called()
-	return &localConfig.ConfigContent{CliId: "134kh"}, nil
+	return &localConfig.LocalConfig{Token: "134kh"}, nil
 }
diff --git a/cmd/publish/main_test.go b/cmd/publish/main_test.go
index 0d91c30..b4453e2 100644
--- a/cmd/publish/main_test.go
+++ b/cmd/publish/main_test.go
@@ -17,9 +17,9 @@ type LocalConfigMock struct {
 	mock.Mock
 }
 
-func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.ConfigContent, error) {
+func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.LocalConfig, error) {
 	lc.Called()
-	return &localConfig.ConfigContent{CliId: "cli_id"}, nil
+	return &localConfig.LocalConfig{Token: "token"}, nil
 }
 
 type MessagerMock struct {
@@ -57,6 +57,15 @@ func (p *PublishClientMock) PublishPolicies(policiesConfiguration files.UnknownS
 	return args.Get(0).(*cliClient.PublishFailedResponse), args.Error(1)
 }
 
+type TokenClientMock struct {
+	mock.Mock
+}
+
+func (t *TokenClientMock) CreateToken() (*cliClient.CreateTokenResponse, error) {
+	args := t.Called()
+	return args.Get(0).(*cliClient.CreateTokenResponse), args.Error(1)
+}
+
 func TestPublishCommand(t *testing.T) {
 	localConfigMock := &LocalConfigMock{}
 	localConfigMock.On("GetLocalConfiguration")
@@ -84,19 +93,19 @@ func TestPublishCommand(t *testing.T) {
 	testPublishCommandFailedSchema(t, ctx, publishClientMock, localConfigContent)
 }
 
-func testPublishCommandSuccess(t *testing.T, ctx *PublishCommandContext, publishClientMock *PublishClientMock, localConfigContent *localConfig.ConfigContent) {
+func testPublishCommandSuccess(t *testing.T, ctx *PublishCommandContext, publishClientMock *PublishClientMock, localConfigContent *localConfig.LocalConfig) {
 	publishClientMock.On("PublishPolicies", mock.Anything, mock.Anything).Return(&cliClient.PublishFailedResponse{}, nil).Once()
 	_, err := publish(ctx, "../../internal/fixtures/policyAsCode/valid-schema.yaml", localConfigContent)
 	assert.Equal(t, nil, err)
 }
 
-func testPublishCommandFailedYaml(t *testing.T, ctx *PublishCommandContext, localConfigContent *localConfig.ConfigContent) {
+func testPublishCommandFailedYaml(t *testing.T, ctx *PublishCommandContext, localConfigContent *localConfig.LocalConfig) {
 	_, err := publish(ctx, "../../internal/fixtures/policyAsCode/invalid-yaml.yaml", localConfigContent)
 	assert.NotEqual(t, nil, err)
 	assert.Equal(t, "yaml: line 2: did not find expected key", err.Error())
 }
 
-func testPublishCommandFailedSchema(t *testing.T, ctx *PublishCommandContext, publishClientMock *PublishClientMock, localConfigContent *localConfig.ConfigContent) {
+func testPublishCommandFailedSchema(t *testing.T, ctx *PublishCommandContext, publishClientMock *PublishClientMock, localConfigContent *localConfig.LocalConfig) {
 	publishFailedPayloadMock := []string{"first error", "second error"}
 	errMessage := strings.Join(publishFailedPayloadMock, ",")
 	publishFailedResponseMock := &cliClient.PublishFailedResponse{
diff --git a/cmd/test/main.go b/cmd/test/main.go
index f8f1197..8d8aa2c 100644
--- a/cmd/test/main.go
+++ b/cmd/test/main.go
@@ -99,7 +99,7 @@ type Reader interface {
 }
 
 type LocalConfig interface {
-	GetLocalConfiguration() (*localConfig.ConfigContent, error)
+	GetLocalConfiguration() (*localConfig.LocalConfig, error)
 }
 
 var ViolationsFoundError = errors.New("")
@@ -195,7 +195,7 @@ func New(ctx *TestCommandContext) *cobra.Command {
 				return err
 			}
 
-			evaluationPrerunData, err := ctx.CliClient.RequestEvaluationPrerunData(localConfigContent.CliId)
+			evaluationPrerunData, err := ctx.CliClient.RequestEvaluationPrerunData(localConfigContent.Token)
 
 			if err != nil {
 				return err
@@ -233,7 +233,7 @@ func (flags *TestCommandFlags) AddFlags(cmd *cobra.Command) {
 	cmd.Flags().BoolVarP(&flags.IgnoreMissingSchemas, "ignore-missing-schemas", "", false, "Ignore missing schemas when executing schema validation step")
 }
 
-func GenerateTestCommandData(testCommandFlags *TestCommandFlags, localConfigContent *localConfig.ConfigContent, evaluationPrerunDataResp *cliClient.EvaluationPrerunDataResponse) (*TestCommandData, error) {
+func GenerateTestCommandData(testCommandFlags *TestCommandFlags, localConfigContent *localConfig.LocalConfig, evaluationPrerunDataResp *cliClient.EvaluationPrerunDataResponse) (*TestCommandData, error) {
 	k8sVersion := testCommandFlags.K8sVersion
 	if k8sVersion == "" {
 		k8sVersion = localConfigContent.SchemaVersion
@@ -253,7 +253,7 @@ func GenerateTestCommandData(testCommandFlags *TestCommandFlags, localConfigCont
 		OnlyK8sFiles:         testCommandFlags.OnlyK8sFiles,
 		Policy:               policy,
 		SchemaLocations:      testCommandFlags.SchemaLocations,
-		Token:                localConfigContent.CliId,
+		Token:                localConfigContent.Token,
 	}
 
 	return testCommandOptions, nil
diff --git a/cmd/test/main_test.go b/cmd/test/main_test.go
index cd4bfb5..ed660b7 100644
--- a/cmd/test/main_test.go
+++ b/cmd/test/main_test.go
@@ -117,9 +117,9 @@ type LocalConfigMock struct {
 	mock.Mock
 }
 
-func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.ConfigContent, error) {
+func (lc *LocalConfigMock) GetLocalConfiguration() (*localConfig.LocalConfig, error) {
 	lc.Called()
-	return &localConfig.ConfigContent{CliId: "134kh"}, nil
+	return &localConfig.LocalConfig{Token: "134kh"}, nil
 }
 
 func TestTestCommand(t *testing.T) {
@@ -183,7 +183,7 @@ func TestTestCommand(t *testing.T) {
 	readerMock.On("FilterFiles", mock.Anything).Return([]string{"file/path"}, nil)
 
 	localConfigMock := &LocalConfigMock{}
-	localConfigMock.On("GetLocalConfiguration").Return(&localConfig.ConfigContent{CliId: "134kh"}, nil)
+	localConfigMock.On("GetLocalConfiguration").Return(&localConfig.LocalConfig{Token: "134kh"}, nil)
 
 	ctx := &TestCommandContext{
 		K8sValidator: k8sValidatorMock,