Merge pull request #7949 from markcowl/import-fix

Fix persistence and cache issues with import-context cmdlet

Author: markcowl

src/ResourceManager/Profile/Commands.Profile.Test/ProfileCmdletTests.cs

@@ -52,7 +52,9 @@ public ProfileCmdletTests(ITestOutputHelper output)
         public void SelectAzureProfileInMemory()
         {
             var profile = new AzureRmProfile { DefaultContext = new AzureContext() };
-            profile.EnvironmentTable.Add("foo", new AzureEnvironment(AzureEnvironment.PublicEnvironments.Values.FirstOrDefault()));
+            var env = new AzureEnvironment(AzureEnvironment.PublicEnvironments.Values.FirstOrDefault());
+            env.Name = "foo";
+            profile.EnvironmentTable.Add("foo", env);
             ImportAzureRMContextCommand cmdlt = new ImportAzureRMContextCommand();
             // Setup
             cmdlt.AzureContext = profile;

src/ResourceManager/Profile/Commands.Profile/Common/AzureContextModificationCmdlet.cs

@@ -47,6 +47,18 @@ protected virtual void ModifyContext(Action<AzureRmProfile, RMProfileClient> con
             }
         }
 
+        /// <summary>
+        /// Modify the Profile according to the selected scope for thsi invocation
+        /// </summary>
+        /// <param name="profileAction">The action to take over the given profile</param>
+        protected virtual void ModifyProfile(Action<IProfileOperations> profileAction)
+        {
+            using (var profile = GetDefaultProfile())
+            {
+                profileAction(profile);
+            }
+        }
+
         /// <summary>
         /// Get the default profile for the current cmdlet invocation
         /// </summary>

src/ResourceManager/Profile/Commands.Profile/Context/ImportAzureRMContext.cs

@@ -15,14 +15,17 @@
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Azure.Commands.Common.Authentication.Models;
+using Microsoft.Azure.Commands.Common.Authentication.ResourceManager;
 using Microsoft.Azure.Commands.Profile.Common;
 using Microsoft.Azure.Commands.Profile.Models;
 // TODO: Remove IfDef
 #if NETSTANDARD
+using Microsoft.Azure.Commands.Common.Authentication.Core;
 using Microsoft.Azure.Commands.Profile.Models.Core;
 #endif
 using Microsoft.Azure.Commands.Profile.Properties;
 using System;
+using System.Linq;
 using System.Management.Automation;
 
 namespace Microsoft.Azure.Commands.Profile
@@ -50,6 +53,55 @@ protected override void BeginProcessing()
             // Do not access the DefaultContext when loading a profile
         }
 
+        void CopyProfile(AzureRmProfile source, IProfileOperations target)
+        {
+            if (source == null || target == null)
+            {
+                return;
+            }
+
+            foreach (var environment in source.Environments)
+            {
+                IAzureEnvironment merged;
+                target.TrySetEnvironment(environment, out merged);
+            }
+
+            foreach (var context in source.Contexts)
+            {
+                target.TrySetContext(context.Key, context.Value);
+            }
+
+            if (!string.IsNullOrWhiteSpace(source.DefaultContextKey))
+            {
+                target.TrySetDefaultContext(source.DefaultContextKey);
+            }
+
+            AzureRmProfileProvider.Instance.SetTokenCacheForProfile(target.ToProfile());
+            EnsureProtectedCache(target, source.DefaultContext?.TokenCache?.CacheData);
+        }
+
+        void EnsureProtectedCache(IProfileOperations profile, byte[] cacheData)
+        {
+            if (profile == null || cacheData == null)
+            {
+                return;
+            }
+
+            AzureRmAutosaveProfile autosave = profile as AzureRmAutosaveProfile;
+            var protectedcache = AzureSession.Instance.TokenCache as ProtectedFileTokenCache;
+            if (autosave != null && protectedcache == null && cacheData.Any())
+            {
+                try
+                {
+                    var cache = new ProtectedFileTokenCache(cacheData, AzureSession.Instance.DataStore);
+                }
+                catch
+                {
+                    WriteWarning(Resources.ImportAuthenticationFailure);
+                }
+            }
+        }
+
         public override void ExecuteCmdlet()
         {
             bool executionComplete = false;
@@ -65,11 +117,9 @@ public override void ExecuteCmdlet()
                             Path));
                     }
 
-                    ModifyContext((profile, client) =>
+                    ModifyProfile((profile) =>
                     {
-                        var newProfile = new AzureRmProfile(Path);
-                        profile.TryCopyProfile(newProfile);
-                        AzureRmProfileProvider.Instance.SetTokenCacheForProfile(newProfile);
+                        CopyProfile(new AzureRmProfile(Path), profile);
                         executionComplete = true;
                     });
                 });
@@ -78,10 +128,9 @@ public override void ExecuteCmdlet()
             {
                 ConfirmAction(Resources.ProcessImportContextFromObject, Resources.ImportContextTarget, () =>
                 {
-                    ModifyContext((profile, client) =>
+                    ModifyProfile((profile) =>
                     {
-                        profile.TryCopyProfile(AzureContext);
-                        AzureRmProfileProvider.Instance.SetTokenCacheForProfile(AzureContext);
+                        CopyProfile(AzureContext, profile);
                         executionComplete = true;
                     });
                 });

src/ResourceManager/Profile/Commands.Profile/Properties/Resources.Designer.cs

@@ -444,4 +444,7 @@
   <data name="ShouldRemoveModule" xml:space="preserve">
     <value>Removing module '{0}' from your machine</value>
   </data>
+  <data name="ImportAuthenticationFailure" xml:space="preserve">
+    <value>Unable to import authentication information into the global cache.  Please try executing the command again.</value>
+  </data>
 </root>