From 9d7ea7483cbc4b9acb6d7885b3c051b7a0dbfc41 Mon Sep 17 00:00:00 2001
From: Steve Kriss <steve@heptio.com>
Date: Thu, 23 Aug 2018 16:24:44 -0700
Subject: [PATCH] azure: support different RGs/storage accounts per backup
 location

Signed-off-by: Steve Kriss <steve@heptio.com>
---
 pkg/cloudprovider/azure/block_store.go  | 34 +++++++--------
 pkg/cloudprovider/azure/object_store.go | 57 +++++++++++++++++++++++--
 2 files changed, 69 insertions(+), 22 deletions(-)

diff --git a/pkg/cloudprovider/azure/block_store.go b/pkg/cloudprovider/azure/block_store.go
index dd7d34d7f..feecfb928 100644
--- a/pkg/cloudprovider/azure/block_store.go
+++ b/pkg/cloudprovider/azure/block_store.go
@@ -40,16 +40,14 @@ import (
 )
 
 const (
-	azureClientIDKey         = "AZURE_CLIENT_ID"
-	azureClientSecretKey     = "AZURE_CLIENT_SECRET"
-	azureSubscriptionIDKey   = "AZURE_SUBSCRIPTION_ID"
-	azureTenantIDKey         = "AZURE_TENANT_ID"
-	azureStorageAccountIDKey = "AZURE_STORAGE_ACCOUNT_ID"
-	azureStorageKeyKey       = "AZURE_STORAGE_KEY"
-	azureResourceGroupKey    = "AZURE_RESOURCE_GROUP"
-	apiTimeoutKey            = "apiTimeout"
-	snapshotsResource        = "snapshots"
-	disksResource            = "disks"
+	azureTenantIDKey       = "AZURE_TENANT_ID"
+	azureSubscriptionIDKey = "AZURE_SUBSCRIPTION_ID"
+	azureClientIDKey       = "AZURE_CLIENT_ID"
+	azureClientSecretKey   = "AZURE_CLIENT_SECRET"
+	azureResourceGroupKey  = "AZURE_RESOURCE_GROUP"
+	apiTimeoutKey          = "apiTimeout"
+	snapshotsResource      = "snapshots"
+	disksResource          = "disks"
 )
 
 type blockStore struct {
@@ -71,15 +69,13 @@ func (si *snapshotIdentifier) String() string {
 	return getComputeResourceName(si.subscription, si.resourceGroup, snapshotsResource, si.name)
 }
 
-func getConfig() map[string]string {
+func getAzureEnvVars() map[string]string {
 	cfg := map[string]string{
-		azureClientIDKey:         "",
-		azureClientSecretKey:     "",
-		azureSubscriptionIDKey:   "",
-		azureTenantIDKey:         "",
-		azureStorageAccountIDKey: "",
-		azureStorageKeyKey:       "",
-		azureResourceGroupKey:    "",
+		azureTenantIDKey:       "",
+		azureSubscriptionIDKey: "",
+		azureClientIDKey:       "",
+		azureClientSecretKey:   "",
+		azureResourceGroupKey:  "",
 	}
 
 	for key := range cfg {
@@ -108,7 +104,7 @@ func (b *blockStore) Init(config map[string]string) error {
 		apiTimeout = 2 * time.Minute
 	}
 
-	cfg := getConfig()
+	cfg := getAzureEnvVars()
 
 	spt, err := helpers.NewServicePrincipalTokenFromCredentials(cfg, azure.PublicCloud.ResourceManagerEndpoint)
 	if err != nil {
diff --git a/pkg/cloudprovider/azure/object_store.go b/pkg/cloudprovider/azure/object_store.go
index 11cf232c0..2673408ac 100644
--- a/pkg/cloudprovider/azure/object_store.go
+++ b/pkg/cloudprovider/azure/object_store.go
@@ -21,7 +21,11 @@ import (
 	"strings"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/arm/examples/helpers"
+	storagemgmt "github.com/Azure/azure-sdk-for-go/services/storage/mgmt/2017-10-01/storage"
 	"github.com/Azure/azure-sdk-for-go/storage"
+	"github.com/Azure/go-autorest/autorest"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 
@@ -37,10 +41,57 @@ func NewObjectStore(logger logrus.FieldLogger) cloudprovider.ObjectStore {
 	return &objectStore{log: logger}
 }
 
-func (o *objectStore) Init(config map[string]string) error {
-	cfg := getConfig()
+func getStorageAccountsClient(envVars map[string]string) (*storagemgmt.AccountsClient, error) {
+	spt, err := helpers.NewServicePrincipalTokenFromCredentials(envVars, azure.PublicCloud.ResourceManagerEndpoint)
+	if err != nil {
+		return nil, errors.Wrap(err, "error creating new service principal token")
+	}
 
-	storageClient, err := storage.NewBasicClient(cfg[azureStorageAccountIDKey], cfg[azureStorageKeyKey])
+	accountsClient := storagemgmt.NewAccountsClient(envVars[azureSubscriptionIDKey])
+	accountsClient.Authorizer = autorest.NewBearerAuthorizer(spt)
+
+	return &accountsClient, nil
+}
+
+func getStorageAccountKey(client *storagemgmt.AccountsClient, resourceGroup, storageAccount string) (string, error) {
+	res, err := client.ListKeys(resourceGroup, storageAccount)
+	if err != nil {
+		return "", errors.WithStack(err)
+	}
+	if res.Keys == nil || len(*res.Keys) == 0 {
+		return "", errors.New("No storage keys found")
+	}
+
+	var storageKey string
+
+	for _, key := range *res.Keys {
+		// uppercase both strings for comparison because the ListKeys call returns e.g. "FULL" but
+		// the storagemgmt.Full constant in the SDK is defined as "Full".
+		if strings.ToUpper(string(key.Permissions)) == strings.ToUpper(string(storagemgmt.Full)) {
+			storageKey = *key.Value
+			break
+		}
+	}
+
+	if storageKey == "" {
+		return "", errors.New("No storage key with Full permissions found")
+	}
+
+	return storageKey, nil
+}
+
+func (o *objectStore) Init(config map[string]string) error {
+	storageAccountsClient, err := getStorageAccountsClient(getAzureEnvVars())
+	if err != nil {
+		return err
+	}
+
+	storageAccountKey, err := getStorageAccountKey(storageAccountsClient, config["resourceGroup"], config["storageAccount"])
+	if err != nil {
+		return err
+	}
+
+	storageClient, err := storage.NewBasicClient(config["storageAccount"], storageAccountKey)
 	if err != nil {
 		return errors.WithStack(err)
 	}