Refactoring alias helper

Author: markcowl

src/ResourceManager/Profile/Commands.Profile/AzureRmAlias/AliasHelper.cs

@@ -12,10 +12,10 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using Microsoft.Azure.Commands.Profile.Properties;
+using Microsoft.Azure.Commands.Common.Authentication;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using System;
 using System.Collections.Generic;
-using System.Collections.ObjectModel;
 using System.IO;
 using System.Linq;
 using System.Management.Automation;
@@ -28,33 +28,47 @@ public class AliasHelper
         private const string STARTALIASIMPORTMARKER = "#Begin Azure PowerShell alias import";
         private const string ENDALIASIMPORTMARKER = "#End Azure PowerShell alias import";
 
+        private IDataStore _dataStore;
+
+        public AliasHelper(IDataStore dataStore)
+        {
+            _dataStore = dataStore;
+            FileUtilities.DataStore = dataStore;
+        }
+
+        public AliasHelper() : this(AzureSession.Instance.DataStore)
+        {
+        }
+
         public static string GetProfilePath(string Scope, SessionState sessionState)
         {
-            object userprofile = "";
+            var userprofile = "";
             if (Scope != null && Scope.Equals("CurrentUser"))
             {
-                userprofile = Path.Combine(sessionState.PSVariable.GetValue("env:USERPROFILE").ToString(), "Documents", "PowerShell", "profile.ps1");
+                var editionType = sessionState.PSVariable.GetValue("PSEdition") as string;
+                var psFolder = string.Equals(editionType, "Desktop", StringComparison.OrdinalIgnoreCase) ? "WindowsPowerShell" : "PowerShell";
+                userprofile = Path.Combine(sessionState.PSVariable.GetValue("env:USERPROFILE").ToString(), "Documents", psFolder, "profile.ps1");
             }
 
             else if (Scope != null && Scope.Equals("LocalMachine"))
             {
                 userprofile = Path.Combine(sessionState.PSVariable.GetValue("PSHOME").ToString(), "profile.ps1");
             }
 
-            return userprofile.ToString();
+            return userprofile;
         }
 
-        public static void RemoveAliasesInProfile(string userprofile, string[] Module, Dictionary<string, object> mapping)
+        public void RemoveAliasesInProfile(string userprofile, string[] Module, Dictionary<string, object> mapping)
         {
             ParseFile(userprofile, Module, out string filecontent, out List<string> modulesToKeep, add: false);
 
             // Add script to enable aliases to profile if there are any modules to keep
             CreateFileEntry(modulesToKeep, filecontent, userprofile, mapping);
         }
 
-        public static void AddAliasesToProfile(string userprofile, string[] Module, Dictionary<string, object> mapping)
+        public void AddAliasesToProfile(string userprofile, string[] Module, Dictionary<string, object> mapping)
         {
-            CreateProfileIfNotExisting(userprofile);
+            CreateProfileIfNotExists(userprofile);
 
             ParseFile(userprofile, Module, out string filecontent, out List<string> modulesToKeep, add: true);
 
@@ -68,15 +82,17 @@ public static void AddAliasesToProfile(string userprofile, string[] Module, Dict
             }
         }
 
-        public static void ParseFile(string userprofile, string[] Module, out string filecontent, out List<string> modulesToKeep, bool add)
+        public void ParseFile(string userprofile, string[] Module, out string filecontent, out List<string> modulesToKeep, bool add)
         {
             filecontent = "";
             modulesToKeep = new List<String>();
             if (add)
             {
                 modulesToKeep = Module?.ToList();
             }
-            string originalText = File.ReadAllText(userprofile.ToString());
+
+            
+            string originalText = _dataStore.ReadFileAsText(userprofile.ToString());
             if (originalText.Contains(STARTALIASIMPORTMARKER))
             {
                 // Add back profile code unrelated to Azure PowerShell aliases
@@ -119,12 +135,12 @@ public static void ParseFile(string userprofile, string[] Module, out string fil
             }
         }
 
-        public static void CreateFileEntry(List<string> modulesToKeep, string filecontent, string userprofile, Dictionary<string, object> mapping)
+        public void CreateFileEntry(List<string> modulesToKeep, string filecontent, string userprofile, Dictionary<string, object> mapping)
         {
             if (modulesToKeep.Count > 0)
             {
                 filecontent += STARTALIASIMPORTMARKER + Environment.NewLine + "Import-Module Az.Profile -ErrorAction SilentlyContinue -ErrorVariable importError" + 
-                    Environment.NewLine + "if (`$importerror.Count -eq 0) { " + Environment.NewLine;
+                    Environment.NewLine + "if ($importerror.Count -eq 0) { " + Environment.NewLine;
 
                 var validModules = new List<string>();
                 foreach (var name in modulesToKeep)
@@ -139,43 +155,34 @@ public static void CreateFileEntry(List<string> modulesToKeep, string fileconten
                 filecontent += "}" + Environment.NewLine + ENDALIASIMPORTMARKER;
             }
 
-            using (System.Management.Automation.PowerShell PowerShellInstance = System.Management.Automation.PowerShell.Create())
+            ExecuteFileActionWithFriendlyError(() =>
             {
-                PowerShellInstance.AddScript("Set-Content -Path '" + userprofile + "' -Value \"" + filecontent + "\"");
-                Collection<PSObject> PSOutput = PowerShellInstance.Invoke();
+                FileUtilities.EnsureDirectoryExists(Path.GetDirectoryName(userprofile));
+                _dataStore.WriteFile(userprofile, filecontent);
+            });
+        }
 
-                // Check for error thrown when LocalMachine profile is accessed from non-admin mode.
-                if (PowerShellInstance.Streams.Error.Count > 0)
+        public void CreateProfileIfNotExists(string userprofile)
+        {
+            ExecuteFileActionWithFriendlyError(() =>
+            {
+                FileUtilities.EnsureDirectoryExists(Path.GetDirectoryName(userprofile));
+                if (!_dataStore.FileExists(userprofile))
                 {
-                    foreach (var error in PowerShellInstance.Streams.Error)
-                    {
-                        if (error.ToString().Contains("Access to the path") && error.ToString().Contains("is denied."))
-                        {
-                            throw new Exception(Resources.AliasImportFailure);
-                        }
-                    }
+                    _dataStore.WriteFile(userprofile, Environment.NewLine);
                 }
-            }
+            });
         }
 
-        public static void CreateProfileIfNotExisting(string userprofile)
+        private static void ExecuteFileActionWithFriendlyError(Action fileAction)
         {
-            using (System.Management.Automation.PowerShell PowerShellInstance = System.Management.Automation.PowerShell.Create())
+            try
             {
-                PowerShellInstance.AddScript("if (!(Test-Path '" + userprofile + "')) { New-Item '" + userprofile + "' -ItemType file -Force }");
-                Collection<PSObject> PSOutput = PowerShellInstance.Invoke();
-
-                // Check for error thrown when LocalMachine profile is accessed from non-admin mode.
-                if (PowerShellInstance.Streams.Error.Count > 0)
-                {
-                    foreach (var error in PowerShellInstance.Streams.Error)
-                    {
-                        if (error.ToString().Contains("Access to the path") && error.ToString().Contains("is denied."))
-                        {
-                            throw new Exception(Resources.AliasImportFailure);
-                        }
-                    }
-                }
+                fileAction();
+            }
+            catch (UnauthorizedAccessException accessException) when (accessException?.Message != null && accessException.Message.ToLower().Contains("denied"))
+            {
+                throw new UnauthorizedAccessException(Properties.Resources.AliasImportFailure, accessException);
             }
         }
     }

src/ResourceManager/Profile/Commands.Profile/AzureRmAlias/DisableAzureRmAlias.cs

@@ -12,6 +12,7 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.ResourceManager.Common;
 using Newtonsoft.Json;
 using System;
@@ -58,9 +59,10 @@ public override void ExecuteCmdlet()
 
             if (Scope != null && (Scope.Equals("CurrentUser") || Scope.Equals("LocalMachine")))
             {
-                if (File.Exists(userprofile.ToString()))
+                if (AzureSession.Instance.DataStore.FileExists(userprofile))
                 {
-                    AliasHelper.RemoveAliasesInProfile(userprofile, Module, mapping);
+                    var helper = new AliasHelper();
+                    helper.RemoveAliasesInProfile(userprofile, Module, mapping);
                 }
             }
         }

src/ResourceManager/Profile/Commands.Profile/AzureRmAlias/EnableAzureRmAlias.cs

@@ -58,7 +58,8 @@ public override void ExecuteCmdlet()
 
             if (Scope != null && (Scope.Equals("CurrentUser") || Scope.Equals("LocalMachine")))
             {
-                AliasHelper.AddAliasesToProfile(userprofile, Module, mapping);
+                var helper = new AliasHelper();
+                helper.AddAliasesToProfile(userprofile, Module, mapping);
             }
         }