Merge pull request Azure#3123 from krkhan/vmbackup_stgaccts_fix

Add support for multiple resource groups in VMBackup

Author: cormacpayne

src/ResourceManager/Compute/Commands.Compute/Extension/AzureVMBackup/AzureVMBackupExtensionProtectedSettings.cs

@@ -12,6 +12,7 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.Commands.Compute.StorageServices;
 using System.Collections.Generic;
 namespace Microsoft.Azure.Commands.Compute.Extension.AzureDiskEncryption
 {
@@ -20,8 +21,12 @@ public class AzureVMBackupBlobSasUris
         public AzureVMBackupBlobSasUris()
         {
             blobSASUri = new List<string>();
+            pageBlobUri = new List<string>();
+            storageCredentialsFactory = new List<StorageCredentialsFactory>();
         }
         public List<string> blobSASUri { get; set; }
+        public List<string> pageBlobUri { get; set; }
+        public List<StorageCredentialsFactory> storageCredentialsFactory { get; set; }
     }
 
     public class AzureVMBackupExtensionProtectedSettings

src/ResourceManager/Compute/Commands.Compute/Extension/AzureVMBackup/AzureVMBackupExtensionUtil.cs

@@ -28,6 +28,8 @@
 using Newtonsoft.Json;
 using System;
 using System.Collections.Generic;
+using System.Linq;
+using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -76,16 +78,16 @@ private string GetBase64Encoding(object obj)
         /// <param name="taskId"></param>
         /// <param name="storageCredentialsFactory"></param>
         /// <returns></returns>
-        public List<CloudPageBlob> FindSnapshot(List<string> blobUris, Dictionary<string, string> snapshotQuery, StorageCredentialsFactory storageCredentialsFactory)
+        public List<CloudPageBlob> FindSnapshot(AzureContext azContext, List<string> blobUris, List<StorageCredentialsFactory> storageCredentialsFactory, Dictionary<string, string> snapshotQuery)
         {
             List<CloudPageBlob> snapshots = new List<CloudPageBlob>();
             for (int i = 0; i < blobUris.Count; i++)
             {
                 BlobUri blobUri = null;
                 if (BlobUri.TryParseUri(new Uri(blobUris[i]), out blobUri))
                 {
-                    StorageCredentials sc = storageCredentialsFactory.Create(blobUri);
-                    CloudStorageAccount cloudStorageAccount = new CloudStorageAccount(sc, true);
+                    StorageCredentials sc = storageCredentialsFactory[i].Create(blobUri);
+                    CloudStorageAccount cloudStorageAccount = new CloudStorageAccount(sc, azContext.Environment.GetEndpointSuffix(AzureEnvironment.Endpoint.StorageEndpointSuffix), true);
                     CloudBlobClient blobClient = cloudStorageAccount.CreateCloudBlobClient();
                     CloudBlobContainer blobContainer = blobClient.GetContainerReference(blobUri.BlobContainerName);
                     IEnumerable<IListBlobItem> blobs = blobContainer.ListBlobs(null, true, BlobListingDetails.All);
@@ -119,9 +121,19 @@ public List<CloudPageBlob> FindSnapshot(List<string> blobUris, Dictionary<string
             }
             return snapshots;
         }
+        internal string GetResourceGroupFromId(string id)
+        {
+            var matcher = new Regex("/subscriptions/([^/]+)/resourceGroups/([^/]+)/providers/(\\w+)");
+            var result = matcher.Match(id);
+            if (!result.Success || result.Groups == null || result.Groups.Count < 3)
+            {
+                throw new InvalidOperationException(string.Format("Cannot find resource group name and storage account name from resource identity {0}", id));
+            }
 
+            return result.Groups[2].Value;
+        }
 
-        public AzureVMBackupBlobSasUris GenerateBlobSasUris(List<string> blobUris, CloudPageBlobObjectFactory cloudPageBlobObjectFactory)
+        public AzureVMBackupBlobSasUris GenerateBlobSasUris(List<string> blobUris, AzureContext azContext)
         {
             AzureVMBackupBlobSasUris blobSASUris = new AzureVMBackupBlobSasUris();
 
@@ -131,6 +143,18 @@ public AzureVMBackupBlobSasUris GenerateBlobSasUris(List<string> blobUris, Cloud
                 BlobUri osBlobUri = null;
                 if (BlobUri.TryParseUri(new Uri(blobUri), out osBlobUri))
                 {
+                    StorageManagementClient storageClient = AzureSession.ClientFactory.CreateArmClient<StorageManagementClient>(azContext, AzureEnvironment.Endpoint.ResourceManager);
+
+                    // Need to convert osBlobUri.StorageAccountName into corresponding resource group name
+
+                    var listResponse = storageClient.StorageAccounts.List();
+                    var account = listResponse.First(accTemp => accTemp.Name.Equals(osBlobUri.StorageAccountName, StringComparison.InvariantCultureIgnoreCase));
+                    string resourceGroupName = GetResourceGroupFromId(account.Id);
+
+                    StorageCredentialsFactory storageCredentialsFactory = new StorageCredentialsFactory(resourceGroupName, storageClient, azContext.Subscription);
+
+                    CloudPageBlobObjectFactory cloudPageBlobObjectFactory = new CloudPageBlobObjectFactory(storageCredentialsFactory, TimeSpan.FromMinutes(1));
+
                     CloudPageBlob pageBlob = cloudPageBlobObjectFactory.Create(osBlobUri);
 
                     SharedAccessBlobPolicy sasConstraints = new SharedAccessBlobPolicy();
@@ -139,6 +163,8 @@ public AzureVMBackupBlobSasUris GenerateBlobSasUris(List<string> blobUris, Cloud
 
                     string sasUri = osBlobUri.Uri + pageBlob.GetSharedAccessSignature(sasConstraints);
                     blobSASUris.blobSASUri.Add(sasUri);
+                    blobSASUris.pageBlobUri.Add(blobUri);
+                    blobSASUris.storageCredentialsFactory.Add(storageCredentialsFactory);
                 }
                 else
                 {
@@ -167,10 +193,11 @@ public void RemoveSnapshot(AzureVMBackupConfig vmConfig, string snapshotTag, Vir
             StorageCredentialsFactory storageCredentialsFactory = new StorageCredentialsFactory(vmConfig.ResourceGroupName, storageClient, virtualMachineExtensionBaseCmdlet.DefaultProfile.Context.Subscription);
 
             List<string> blobUris = this.GetDiskBlobUris(virtualMachineResponse.Body);
+            AzureVMBackupBlobSasUris blobSASUris = this.GenerateBlobSasUris(blobUris, virtualMachineExtensionBaseCmdlet.DefaultProfile.Context);
 
             Dictionary<string, string> snapshotQuery = new Dictionary<string, string>();
             snapshotQuery.Add(backupExtensionMetadataName, snapshotTag);
-            List <CloudPageBlob> snapshots = this.FindSnapshot(blobUris, snapshotQuery, storageCredentialsFactory);
+            List<CloudPageBlob> snapshots = this.FindSnapshot(virtualMachineExtensionBaseCmdlet.DefaultProfile.Context, blobSASUris.pageBlobUri, blobSASUris.storageCredentialsFactory, snapshotQuery);
             if (snapshots == null || snapshots.Count == 0)
             {
                 throw new AzureVMBackupException(AzureVMBackupErrorCodes.NoSnapshotFound, "snapshot with the tag not found.");
@@ -200,15 +227,10 @@ public void CreateSnapshotForDisks(AzureVMBackupConfig vmConfig, string snapshot
             var virtualMachine = virtualMachineExtensionBaseCmdlet.ComputeClient.ComputeManagementClient.VirtualMachines.GetWithInstanceView(
                 vmConfig.ResourceGroupName,
                 vmConfig.VMName);
-            StorageManagementClient storageClient = AzureSession.ClientFactory.CreateArmClient<StorageManagementClient>(virtualMachineExtensionBaseCmdlet.DefaultProfile.Context, AzureEnvironment.Endpoint.ResourceManager);
-
-            StorageCredentialsFactory storageCredentialsFactory = new StorageCredentialsFactory(vmConfig.ResourceGroupName, storageClient, virtualMachineExtensionBaseCmdlet.DefaultProfile.Context.Subscription);
-
-            CloudPageBlobObjectFactory cloudPageBlobObjectFactory = new CloudPageBlobObjectFactory(storageCredentialsFactory, TimeSpan.FromMinutes(1));
 
             List<string> vmPageBlobUris = this.GetDiskBlobUris(virtualMachine.Body);
 
-            AzureVMBackupBlobSasUris blobSASUris = this.GenerateBlobSasUris(vmPageBlobUris, cloudPageBlobObjectFactory);
+            AzureVMBackupBlobSasUris blobSASUris = this.GenerateBlobSasUris(vmPageBlobUris, virtualMachineExtensionBaseCmdlet.DefaultProfile.Context);
 
             string taskId = Guid.NewGuid().ToString();
 
@@ -262,7 +284,7 @@ public void CreateSnapshotForDisks(AzureVMBackupConfig vmConfig, string snapshot
             int i = 0;
             for (; i < loopingTimes; i++)
             {
-                List<CloudPageBlob> snapshotsFound = this.FindSnapshot(vmPageBlobUris, snapshotQuery, storageCredentialsFactory);
+                List<CloudPageBlob> snapshotsFound = this.FindSnapshot(virtualMachineExtensionBaseCmdlet.DefaultProfile.Context, blobSASUris.pageBlobUri, blobSASUris.storageCredentialsFactory, snapshotQuery);
                 if (snapshotsFound.Count == vmPageBlobUris.Count)
                 {
                     break;