azkv: enable specifying some auth methods directly, and add cachable user authentication methods

Author: daytime-saxophone

README.rst

@@ -336,6 +336,14 @@ which tries several authentication methods, in this order:
 3. `Managed Identity credentials <https://pkg.go.dev/github.com/Azure/azure-sdk-for-go/sdk/azidentity#ManagedIdentityCredential>`_
 4. `Azure CLI credentials <https://pkg.go.dev/github.com/Azure/azure-sdk-for-go/sdk/azidentity#AzureCLICredential>`_
 
+
+If you want to force a specific method you can override this with the enviornment variable ``SOPS_AZURE_AUTH_METHOD``
+- ``default`` (same as not setting this variable)
+- ``msi``
+- ``azure-cli``
+- ``cached-device-code`` (device code authentication which caches the token in the os keyring)
+- ``cached-browser`` (interactive browser authentication which caches the token in the os keyring)
+
 For example, you can use a Service Principal with the following environment variables:
 
 .. code:: bash

azkv/keysource.go

@@ -9,13 +9,19 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"path/filepath"
 	"regexp"
 	"strings"
 	"time"
 
+	"encoding/json"
+	"os"
+
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	azidentitycache "github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache"
 	"github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys"
 	"github.com/sirupsen/logrus"
 
@@ -25,6 +31,11 @@ import (
 const (
 	// KeyTypeIdentifier is the string used to identify an Azure Key Vault MasterKey.
 	KeyTypeIdentifier = "azure_kv"
+
+	SopsAzureAuthMethodEnv = "SOPS_AZURE_AUTH_METHOD"
+
+	cachedBrowserAuthRecordFileName    = "azure-auth-record-browser.json"
+	cachedDeviceCodeAuthRecordFileName = "azure-auth-record-device-code.json"
 )
 
 var (
@@ -230,7 +241,130 @@ func (key *MasterKey) TypeToIdentifier() string {
 // azidentity.NewDefaultAzureCredential.
 func (key *MasterKey) getTokenCredential() (azcore.TokenCredential, error) {
 	if key.tokenCredential == nil {
-		return azidentity.NewDefaultAzureCredential(nil)
+
+		authMethod := strings.ToLower(os.Getenv(SopsAzureAuthMethodEnv))
+		switch authMethod {
+		case "cached-browser":
+			return cachedInteractiveBrowserCredentials()
+		case "cached-device-code":
+			return cachedDeviceCodeCredentials()
+		case "azure-cli":
+			return azidentity.NewAzureCLICredential(nil)
+		case "msi":
+			return azidentity.NewManagedIdentityCredential(nil)
+		// If "DEFAULT" or not explicitly specified then use the default authentication chain.
+		case "", "default":
+			return azidentity.NewDefaultAzureCredential(nil)
+		default:
+			return nil, fmt.Errorf("Value `%s` is unsupported for environment variable `%s`, to resolve this either leave it unset or use one of `default`/`msi`/`azure-cli`/`cached-browser`/`cached-device-code`", authMethod, SopsAzureAuthMethodEnv)
+		}
 	}
 	return key.tokenCredential, nil
 }
+
+func sopsCacheDir() (string, error) {
+	userCacheDir, err := os.UserCacheDir()
+	if err != nil {
+		return "", err
+	}
+
+	cacheDir := filepath.Join(userCacheDir, "/sops")
+
+	if err = os.MkdirAll(cacheDir, 0o700); err != nil {
+		return "", err
+	}
+
+	return cacheDir, nil
+}
+
+type CachableTokenCredential interface {
+	Authenticate(ctx context.Context, opts *policy.TokenRequestOptions) (azidentity.AuthenticationRecord, error)
+	GetToken(ctx context.Context, opts policy.TokenRequestOptions) (azcore.AccessToken, error)
+}
+
+func cacheStoreRecord(cachePath string, record azidentity.AuthenticationRecord) error {
+	b, err := json.Marshal(record)
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(cachePath, b, 0600)
+}
+
+func cacheLoadRecord(cachePath string) (azidentity.AuthenticationRecord, error) {
+	var record azidentity.AuthenticationRecord
+
+	b, err := os.ReadFile(cachePath)
+	if err != nil {
+		return record, err
+	}
+
+	err = json.Unmarshal(b, &record)
+	if err != nil {
+		return record, err
+	}
+
+	return record, nil
+}
+
+func cacheTokenCredential(cachePath string, tokenCredentialFn func(cache azidentity.Cache, record azidentity.AuthenticationRecord) (CachableTokenCredential, error)) (azcore.TokenCredential, error) {
+	cache, err := azidentitycache.New(nil)
+	// Errors if persistent caching is not supported by the current runtime
+	if err != nil {
+		return nil, err
+	}
+
+	cachedRecord, cacheLoadErr := cacheLoadRecord(cachePath)
+
+	credential, err := tokenCredentialFn(cache, cachedRecord)
+	if err != nil {
+		return nil, err
+	}
+
+	// If loading the authenticationRecord from the cachePath failed for any reason (validation, file doesn't exist, not encoded using json, etc.)
+	if cacheLoadErr != nil {
+		record, err := credential.Authenticate(context.Background(), nil)
+		if err != nil {
+			return nil, err
+		}
+
+		if err = cacheStoreRecord(cachePath, record); err != nil {
+			return nil, err
+		}
+	}
+
+	return credential, nil
+}
+
+func cachedInteractiveBrowserCredentials() (azcore.TokenCredential, error) {
+	cacheDir, err := sopsCacheDir()
+	if err != nil {
+		return nil, err
+	}
+	return cacheTokenCredential(
+		filepath.Join(cacheDir, cachedBrowserAuthRecordFileName),
+		func(cache azidentity.Cache, record azidentity.AuthenticationRecord) (CachableTokenCredential, error) {
+			return azidentity.NewInteractiveBrowserCredential(&azidentity.InteractiveBrowserCredentialOptions{
+				AuthenticationRecord: record,
+				Cache:                cache,
+			})
+		},
+	)
+}
+
+func cachedDeviceCodeCredentials() (azcore.TokenCredential, error) {
+	cacheDir, err := sopsCacheDir()
+	if err != nil {
+		return nil, err
+	}
+
+	return cacheTokenCredential(
+		filepath.Join(cacheDir, cachedDeviceCodeAuthRecordFileName),
+		func(cache azidentity.Cache, record azidentity.AuthenticationRecord) (CachableTokenCredential, error) {
+			return azidentity.NewDeviceCodeCredential(&azidentity.DeviceCodeCredentialOptions{
+				AuthenticationRecord: record,
+				Cache:                cache,
+			})
+		},
+	)
+}