lazysecrets/internal/provider/databricks/databricks_test.go

package databricks

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"

	dbconfig "github.com/databricks/databricks-sdk-go/config"
	"github.com/databricks/databricks-sdk-go/service/workspace"

	"git.kebcloud.no/keb/lazysecrets/internal/provider"
)

func TestReadProfiles(t *testing.T) {
	path := filepath.Join(t.TempDir(), "databrickscfg")
	cfg := `[DEFAULT]
host      = https://adb-1.1.azuredatabricks.net
auth_type = azure-cli

[INFRA-DEV]
host      = adb-2.2.azuredatabricks.net
auth_type = azure-cli

[ACCOUNT]
host       = https://accounts.azuredatabricks.net
auth_type  = azure-cli

[acct-by-id]
host       = https://example.cloud.databricks.com
account_id = 123

[oauth]
host      = https://adb-3.3.azuredatabricks.net
auth_type = databricks-cli

[no-host]
token = x
`
	if err := os.WriteFile(path, []byte(cfg), 0o600); err != nil {
		t.Fatal(err)
	}
	got, err := readProfiles(path)
	if err != nil {
		t.Fatal(err)
	}
	want := []profile{
		{name: "DEFAULT", host: "https://adb-1.1.azuredatabricks.net", authType: "azure-cli"},
		{name: "INFRA-DEV", host: "adb-2.2.azuredatabricks.net", authType: "azure-cli"},
		{name: "oauth", host: "https://adb-3.3.azuredatabricks.net", authType: "databricks-cli"},
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("profiles = %+v\nwant       %+v", got, want)
	}

	if got, err := readProfiles(filepath.Join(t.TempDir(), "missing")); err != nil || got != nil {
		t.Errorf("missing file: %v, %v", got, err)
	}
}

func TestLoginCommand(t *testing.T) {
	for authType, want := range map[string][]string{
		"azure-cli":      {"az", "login"},
		"databricks-cli": {"databricks", "auth", "login", "--profile", "p"},
		"":               {"databricks", "auth", "login", "--profile", "p"},
		"pat":            nil,
	} {
		if got := loginCommand(profile{name: "p", authType: authType}); !reflect.DeepEqual(got, want) {
			t.Errorf("%q: %v, want %v", authType, got, want)
		}
	}
}

// fakeKeyVault stands in for the Azure provider.
type fakeKeyVault struct {
	secrets map[string]string // name → value, for vault URL kvURL
	err     error
	calls   int
}

const kvURL = "https://kv-app-prod.vault.azure.net/"

func (k *fakeKeyVault) Objects(_ context.Context, c provider.Container, _ provider.ObjectType) ([]provider.Object, error) {
	k.calls++
	if k.err != nil {
		return nil, k.err
	}
	if c.Detail != kvURL {
		return nil, errors.New("wrong vault " + c.Detail)
	}
	var out []provider.Object
	for name := range k.secrets {
		out = append(out, provider.Object{Name: name})
	}
	return out, nil
}

func (k *fakeKeyVault) Value(_ context.Context, o provider.Object) (provider.SecretValue, error) {
	k.calls++
	if k.err != nil {
		return provider.SecretValue{}, k.err
	}
	if o.Container.Detail != kvURL {
		return provider.SecretValue{}, errors.New("wrong vault " + o.Container.Detail)
	}
	return provider.NewSecretValue([]byte(k.secrets[o.Name])), nil
}

// workspaceAPI serves the secrets API. The "akv" scope behaves like an Azure
// Key Vault-backed scope whose secrets Databricks refuses to read.
func workspaceAPI(t *testing.T, status int) *httptest.Server {
	t.Helper()
	write := func(w http.ResponseWriter, code int, v any) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(code)
		_ = json.NewEncoder(w).Encode(v)
	}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if status != 0 {
			write(w, status, map[string]any{"error_code": "UNAUTHENTICATED", "message": "Token is expired"})
			return
		}
		q := r.URL.Query()
		switch {
		case r.URL.Path == "/api/2.0/secrets/scopes/list":
			write(w, 200, map[string]any{"scopes": []any{
				map[string]any{"name": "app", "backend_type": "DATABRICKS"},
				map[string]any{"name": "akv", "backend_type": "AZURE_KEYVAULT",
					"keyvault_metadata": map[string]any{"dns_name": kvURL, "resource_id": "/x/kv-app-prod"}},
			}})
		case r.URL.Path == "/api/2.0/secrets/list" && q.Get("scope") == "app":
			write(w, 200, map[string]any{"secrets": []any{
				map[string]any{"key": "storage-key", "last_updated_timestamp": 1767225600000},
				map[string]any{"key": "client-secret"},
			}})
		case r.URL.Path == "/api/2.0/secrets/list" && q.Get("scope") == "locked":
			write(w, 403, map[string]any{"error_code": "PERMISSION_DENIED", "message": "User does not have READ permission on scope locked"})
		case r.URL.Path == "/api/2.0/secrets/get" && q.Get("scope") == "app":
			write(w, 200, map[string]any{"key": q.Get("key"),
				"value": base64.StdEncoding.EncodeToString([]byte("value-of-" + q.Get("key")))})
		case strings.HasPrefix(r.URL.Path, "/api/2.0/secrets/") && q.Get("scope") == "akv":
			write(w, 400, map[string]any{"error_code": "BAD_REQUEST",
				"message": "Reading secrets from Azure KeyVault-backed scopes is not supported"})
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(srv.Close)
	return srv
}

func newTestProvider(t *testing.T, srv *httptest.Server, kv KeyVault) *Provider {
	t.Helper()
	return &Provider{
		profiles: []profile{{name: "prod", host: srv.URL, authType: "azure-cli"}},
		keyVault: kv,
		newClient: func(profile) (*workspace.SecretsAPI, error) {
			return secretsClient(&dbconfig.Config{Host: srv.URL, Token: "token"})
		},
		clients: map[string]*workspace.SecretsAPI{},
		vaults:  map[string]string{},
	}
}

func scopes(t *testing.T, p *Provider) map[string]provider.Container {
	t.Helper()
	srcs, _ := p.Sources(context.Background())
	cs, err := p.Containers(context.Background(), srcs[0])
	if err != nil {
		t.Fatal(err)
	}
	out := map[string]provider.Container{}
	for _, c := range cs {
		out[c.Name] = c
	}
	return out
}

func TestScopesSecretsAndValues(t *testing.T) {
	p := newTestProvider(t, workspaceAPI(t, 0), nil)
	sc := scopes(t, p)
	if sc["app"].Detail != "Databricks-backed" || sc["akv"].Detail != "Azure Key Vault-backed: "+kvURL {
		t.Errorf("scope details = %q / %q", sc["app"].Detail, sc["akv"].Detail)
	}

	objs, err := p.Objects(context.Background(), sc["app"], provider.TypeSecret)
	if err != nil {
		t.Fatal(err)
	}
	if len(objs) != 2 || objs[0].Name != "client-secret" || objs[1].Name != "storage-key" {
		t.Fatalf("objects = %+v", objs)
	}
	if objs[1].Updated.IsZero() || objs[1].Reference != "{{secrets/app/storage-key}}" {
		t.Errorf("metadata = %+v", objs[1])
	}

	v, err := p.Value(context.Background(), objs[1])
	if err != nil || v.Reveal() != "value-of-storage-key" {
		t.Errorf("value = %q, %v", v.Reveal(), err)
	}
}

func TestAKVScopeFallsBackToKeyVault(t *testing.T) {
	kv := &fakeKeyVault{secrets: map[string]string{"db-password": "from-the-vault"}}
	p := newTestProvider(t, workspaceAPI(t, 0), kv)
	akv := scopes(t, p)["akv"]

	objs, err := p.Objects(context.Background(), akv, provider.TypeSecret)
	if err != nil {
		t.Fatal(err)
	}
	if len(objs) != 1 || objs[0].Name != "db-password" || objs[0].Reference != "{{secrets/akv/db-password}}" {
		t.Fatalf("objects = %+v", objs)
	}
	v, err := p.Value(context.Background(), objs[0])
	if err != nil || v.Reveal() != "from-the-vault" {
		t.Errorf("value = %q, %v", v.Reveal(), err)
	}
}

func TestAKVFallbackForRememberedScope(t *testing.T) {
	// A favourite from an earlier session: scopes not listed yet.
	kv := &fakeKeyVault{secrets: map[string]string{"db-password": "from-the-vault"}}
	p := newTestProvider(t, workspaceAPI(t, 0), kv)
	akv := provider.Container{
		Source: provider.Source{Kind: provider.KindDatabricks, ID: "prod"},
		ID:     "prod/akv", Name: "akv", Detail: "Azure Key Vault-backed: " + kvURL,
	}
	v, err := p.Value(context.Background(), provider.Object{Container: akv, Name: "db-password"})
	if err != nil || v.Reveal() != "from-the-vault" {
		t.Errorf("value = %q, %v", v.Reveal(), err)
	}
}

func TestAKVFallbackFailureReportsBoth(t *testing.T) {
	kvAuth := &provider.AuthError{Kind: provider.KindAzure, Login: []string{"az", "login"}, Err: errors.New("expired")}
	p := newTestProvider(t, workspaceAPI(t, 0), &fakeKeyVault{err: kvAuth})
	akv := scopes(t, p)["akv"]

	_, err := p.Value(context.Background(), provider.Object{Container: akv, Name: "db-password"})
	if err == nil || !strings.Contains(err.Error(), "not supported") || !strings.Contains(err.Error(), "Key Vault fallback") {
		t.Fatalf("err = %v", err)
	}
	var auth *provider.AuthError
	if !errors.As(err, &auth) {
		t.Error("Key Vault auth error not reachable for the sign-in prompt")
	}
}

func TestNoFallbackForDatabricksBackedScopes(t *testing.T) {
	kv := &fakeKeyVault{}
	p := newTestProvider(t, workspaceAPI(t, 0), kv)
	app := scopes(t, p)["app"]
	if _, err := p.Value(context.Background(), provider.Object{Container: app, Name: "missing"}); err != nil {
		// app/get returns a value for any key; just make sure the vault wasn't asked
		t.Fatal(err)
	}
	if kv.calls != 0 {
		t.Errorf("Key Vault consulted for a Databricks-backed scope")
	}
}

func TestUnauthorizedIsAuthError(t *testing.T) {
	p := newTestProvider(t, workspaceAPI(t, 401), nil)
	srcs, _ := p.Sources(context.Background())
	_, err := p.Containers(context.Background(), srcs[0])
	var auth *provider.AuthError
	if !errors.As(err, &auth) {
		t.Fatalf("err = %v, want AuthError", err)
	}
	if auth.Kind != provider.KindDatabricks || !reflect.DeepEqual(auth.Login, []string{"az", "login"}) {
		t.Errorf("auth = %+v", auth)
	}
}

func TestClassify(t *testing.T) {
	pr := profile{name: "p", authType: "databricks-cli"}
	var auth *provider.AuthError
	if !errors.As(classify(pr, errors.New("default auth: azure-cli: cannot get access token: run az login")), &auth) {
		t.Error("token failure not classified as auth")
	}
	if errors.As(classify(pr, errors.New("dial tcp: no such host")), &auth) {
		t.Error("network error classified as auth")
	}
	if errors.As(classify(pr, context.Canceled), &auth) {
		t.Error("cancellation classified as auth")
	}
}

func TestDeniedSaysWhatToCheck(t *testing.T) {
	p := newTestProvider(t, workspaceAPI(t, 0), nil)
	_, err := p.Objects(context.Background(), provider.Container{Source: provider.Source{ID: "prod"}, Name: "locked"}, provider.TypeSecret)
	var d *provider.Denied
	// The test profile signs in through az: that's the first check.
	if !errors.As(err, &d) || !strings.Contains(err.Error(), "READ permission on scope locked\n\nCheck:\n• az is signed in as the account with access: profile prod uses your az login") ||
		!strings.Contains(err.Error(), "\n• you have access to the workspace") {
		t.Errorf("err %q", err)
	}
	if got := deniedChecks(profile{authType: "pat"}, "denied"); strings.Contains(got[0], "az") {
		t.Errorf("a token profile's checks %q", got)
	}
	if got := deniedChecks(profile{}, "Source IP address: 1.2.3.4 is blocked by Databricks IP ACL for workspace: 1"); !strings.Contains(got[0], "network") {
		t.Errorf("IP ACL checks %q", got)
	}

	// Both refused: one message, then both lists of checks.
	err = combine(&provider.Denied{Err: errors.New("dbx"), Checks: []string{"a", "b"}},
		&provider.Denied{Err: errors.New("kv"), Checks: []string{"b", "c"}})
	if got := err.Error(); got != "dbx; Key Vault fallback: kv\n\nCheck:\n• a\n• b\n• c" {
		t.Errorf("combined %q", got)
	}
}