From 712e2919f5d6d1b7df396bb067a2e3c853d0408d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Sun, 25 Aug 2024 23:53:44 +0200 Subject: [PATCH] add logout cmd --- cmd/auth/auth.go | 1 + cmd/auth/logout.go | 102 ++++++++++++++++++++++++++++++++++++ cmd/auth/logout_test.go | 113 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 216 insertions(+) create mode 100644 cmd/auth/logout.go create mode 100644 cmd/auth/logout_test.go diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index ceceae25..57b7b8a4 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -31,6 +31,7 @@ GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html`, cmd.AddCommand(newProfilesCommand()) cmd.AddCommand(newTokenCommand(&perisistentAuth)) cmd.AddCommand(newDescribeCommand()) + cmd.AddCommand(newLogoutCommand()) return cmd } diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go new file mode 100644 index 00000000..1e51645e --- /dev/null +++ b/cmd/auth/logout.go @@ -0,0 +1,102 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "io/fs" + + "github.com/databricks/cli/libs/auth/cache" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" +) + +type Logout struct { + Profile string + File config.File + Cache cache.TokenCache +} + +func (l *Logout) load(ctx context.Context, profileName string) error { + l.Profile = profileName + l.Cache = cache.GetTokenCache(ctx) + iniFile, err := profile.DefaultProfiler.Get(ctx) + if errors.Is(err, fs.ErrNotExist) { + return err + } else if err != nil { + return fmt.Errorf("cannot parse config file: %w", err) + } + l.File = *iniFile + return nil +} + +func (l *Logout) getSetionMap() (map[string]string, error) { + section, err := l.File.GetSection(l.Profile) + if err != nil { + return map[string]string{}, fmt.Errorf("profile does not exist in config file: %w", err) + } + return section.KeysHash(), nil +} + +// clear token from ~/.databricks/token-cache.json +func (l *Logout) clearTokenCache(key string) error { + return l.Cache.DeleteKey(key) +} + +// Overrewrite profile to .databrickscfg without fields marked as sensitive +// Other attributes are preserved. +func (l *Logout) clearConfigFile(ctx context.Context, sectionMap map[string]string) error { + return databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: l.File.Path(), + Profile: l.Profile, + Host: sectionMap["host"], + ClusterID: sectionMap["cluster_id"], + WarehouseID: sectionMap["warehouse_id"], + ServerlessComputeID: sectionMap["serverless_compute_id"], + AccountID: sectionMap["account_id"], + Username: sectionMap["username"], + GoogleServiceAccount: sectionMap["google_service_account"], + AzureResourceID: sectionMap["azure_workspace_resource_id"], + AzureClientID: sectionMap["azure_client_id"], + AzureTenantID: sectionMap["azure_tenant_id"], + AzureEnvironment: sectionMap["azure_environment"], + AzureLoginAppID: sectionMap["azure_login_app_id"], + ClientID: sectionMap["client_id"], + AuthType: sectionMap["auth_type"], + }) +} + +func newLogoutCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "logout [PROFILE]", + Short: "Logout from specified profile", + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + var profileName string + if len(args) < 1 { + profileName = cmd.Flag("profile").Value.String() + } else { + profileName = args[0] + } + logout := &Logout{} + logout.load(ctx, profileName) + sectionMap, err := logout.getSetionMap() + if err != nil { + return err + } + if err := logout.clearTokenCache(sectionMap["host"]); err != nil { + return err + } + if err := logout.clearConfigFile(ctx, sectionMap); err != nil { + return err + } + cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully logged out", profileName)) + return nil + } + return cmd +} diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go new file mode 100644 index 00000000..e3344489 --- /dev/null +++ b/cmd/auth/logout_test.go @@ -0,0 +1,113 @@ +package auth + +import ( + "context" + "fmt" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" + + "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/databricks-sdk-go/config" +) + +func TestLogout_ClearConfigFile(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "databrickscfg") + + err := databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: path, + Profile: "abc", + Host: "https://foo", + Token: "xyz", + }) + require.NoError(t, err) + iniFile, err := config.LoadFile(path) + require.NoError(t, err) + logout := &Logout{ + Profile: "abc", + File: *iniFile, + } + section, err := logout.File.GetSection("abc") + assert.NoError(t, err) + sectionMap := section.KeysHash() + err = logout.clearConfigFile(ctx, sectionMap) + assert.NoError(t, err) + + iniFile, err = config.LoadFile(path) + require.NoError(t, err) + assert.Len(t, iniFile.Sections(), 2) + assert.True(t, iniFile.HasSection("DEFAULT")) + assert.True(t, iniFile.HasSection("abc")) + + abc, err := iniFile.GetSection("abc") + assert.NoError(t, err) + raw := abc.KeysHash() + assert.Len(t, raw, 1) + assert.Equal(t, "https://foo", raw["host"]) +} + +type tokenCacheMock struct { + store func(key string, t *oauth2.Token) error + lookup func(key string) (*oauth2.Token, error) + deleteKey func(key string) error +} + +func (m *tokenCacheMock) Store(key string, t *oauth2.Token) error { + if m.store == nil { + panic("no store mock") + } + return m.store(key, t) +} + +func (m *tokenCacheMock) Lookup(key string) (*oauth2.Token, error) { + if m.lookup == nil { + panic("no lookup mock") + } + return m.lookup(key) +} + +func (m *tokenCacheMock) DeleteKey(key string) error { + if m.deleteKey == nil { + panic("no deleteKey mock") + } + return m.deleteKey(key) +} + +func TestLogout_ClearTokenCache(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "databrickscfg") + + err := databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: path, + Profile: "abc", + Host: "https://foo", + AuthType: "databricks-cli", + }) + require.NoError(t, err) + iniFile, err := config.LoadFile(path) + require.NoError(t, err) + logout := &Logout{ + Profile: "abc", + File: *iniFile, + Cache: &tokenCacheMock{ + deleteKey: func(key string) error { + assert.Equal(t, "https://foo", key) + return nil + }, + lookup: func(key string) (*oauth2.Token, error) { + assert.Equal(t, "https://foo", key) + return &oauth2.Token{}, fmt.Errorf("No token found") + }, + }, + } + sectionMap, err := logout.getSetionMap() + assert.NoError(t, err) + err = logout.clearTokenCache(sectionMap["host"]) + assert.NoError(t, err) + _, err = logout.Cache.Lookup(sectionMap["host"]) + assert.Error(t, err) +}