Change update process to modify and restore execution policy before and after update.

Change update process to wait for the client to close instead of a static wait.

Author: Alexandru Macocian

Daybreak/Configuration/ProjectConfiguration.cs

@@ -39,6 +39,7 @@ public static void RegisterLifetimeServices(IApplicationLifetimeProducer applica
 
             applicationLifetimeProducer.RegisterService<ILoggingDatabase>();
             applicationLifetimeProducer.RegisterService<IScreenshotProvider>();
+            applicationLifetimeProducer.RegisterService<IApplicationUpdater>();
         }
         public static void RegisterViews(IViewProducer viewProducer)
         {

Daybreak/Daybreak.csproj

@@ -9,7 +9,7 @@
     <CopyLocalLockFileAssemblies>false</CopyLocalLockFileAssemblies>
     <LangVersion>preview</LangVersion>
     <ApplicationIcon>Daybreak.ico</ApplicationIcon>
-    <Version>0.2.1</Version>
+    <Version>0.2.2</Version>
   </PropertyGroup>
 
   <ItemGroup>

Daybreak/Models/ExecutionPolicies.cs

@@ -0,0 +1,13 @@
+namespace Daybreak.Models
+{
+    public enum ExecutionPolicies
+    {
+        AllSigned,
+        Bypass,
+        Default,
+        RemoteSigned,
+        Restricted,
+        Undefined,
+        Unrestricted
+    }
+}

Daybreak/Services/Updater/ApplicationUpdater.cs

@@ -1,6 +1,7 @@
 using Daybreak.Models;
 using Daybreak.Services.Logging;
 using Daybreak.Utils;
+using Microsoft.Win32;
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
@@ -15,15 +16,21 @@ namespace Daybreak.Services.Updater
 {
     public sealed class ApplicationUpdater : IApplicationUpdater
     {
+        private const string ExecutionPolicyKey = "ExecutionPolicy";
+        private const string UpdatedKey = "Updating";
+        private const string RegistryKey = "Daybreak";
         private const string ExtractAndRunPs1 = "ExtractAndRun.ps1";
         private const string TempFile = "tempfile.zip";
         private const string VersionTag = "{VERSION}";
         private const string InputFileTag = "{INPUTFILE}";
-        private const string OutputPathTag = "{OUTPUTPATh}";
+        private const string OutputPathTag = "{OUTPUTPATH}";
+        private const string ExecutionPolicyTag = "{EXECUTIONPOLICY}";
+        private const string ProcessIdTag = "{PROCESSID}";
         private const string Url = "https://github.com/AlexMacocian/Daybreak/releases/latest";
         private const string DownloadUrl = $"https://github.com/AlexMacocian/Daybreak/releases/download/v{VersionTag}/Daybreakv{VersionTag}.zip";
-        private const string SetExecutionPolicy = $"Set-ExecutionPolicy RemoteSigned -Scope CurrentUser";
-        private const string DelayCommand = "Start-Sleep -m 3000";
+        private const string GetExecutionPolicyCommand = "Get-ExecutionPolicy -Scope CurrentUser";
+        private const string SetExecutionPolicyCommand = $"Set-ExecutionPolicy {ExecutionPolicyTag} -Scope CurrentUser";
+        private const string WaitCommand = $"Wait-Process -Id {ProcessIdTag}";
         private const string ExtractCommandTemplate = $"Expand-Archive -Path '{InputFileTag}' -DestinationPath '{OutputPathTag}' -Force";
         private const string RunClientCommand = @".\Daybreak.exe";
         private const string RemoveTempFile = $"Remove-item {TempFile}";
@@ -90,19 +97,125 @@ public async Task<bool> UpdateAvailable()
             var maybeLatestVersion = await this.GetLatestVersion();
             return maybeLatestVersion.Switch(
                 onSome: latestVersion => string.Compare(version, latestVersion, true) < 0,
-                onNone: () => 
+                onNone: () =>
                 {
                     this.logger.LogWarning("Failed to retrieve latest version");
                     return false;
                 }).ExtractValue();
         }
 
         public void FinalizeUpdate()
+        {
+            var maybeExecutionPolicy = this.RetrieveExecutionPolicy();
+            maybeExecutionPolicy.DoAny(
+                onNone: () =>
+                {
+                    throw new InvalidOperationException("Failed to retrieve execution policy");
+                });
+
+            var executionPolicy = maybeExecutionPolicy.ExtractValue();
+            if (executionPolicy is not ExecutionPolicies.Bypass ||
+                executionPolicy is not ExecutionPolicies.Unrestricted)
+            {
+                this.logger.LogInformation($"Execution policy is set to {executionPolicy}. Setting to {ExecutionPolicies.Bypass}");
+            }
+
+            SaveExecutionPolicyValueToRegistry(executionPolicy);
+            MarkUpdateInRegistry();
+            this.SetExecutionPolicy(ExecutionPolicies.Bypass);
+            this.LaunchExtractor();
+        }
+
+        public void OnStartup()
+        {
+            if (UpdateMarkedInRegistry())
+            {
+                UnmarkUpdateInRegistry();
+                var maybeExecutionPolicy = LoadExecutionPolicyValueFromRegistry();
+                maybeExecutionPolicy.Do(
+                    onSome: policy =>
+                    {
+                        SetExecutionPolicy(policy);
+                    },
+                    onNone: () =>
+                    {
+                        throw new InvalidOperationException("Found update marked in registry but no execution policy");
+                    });
+            }
+        }
+
+        public void OnClosing()
+        {
+        }
+
+        private async Task<Optional<string>> GetLatestVersion()
+        {
+            using var response = await this.httpClient.GetAsync(Url);
+            if (response.IsSuccessStatusCode)
+            {
+                var versionTag = response.RequestMessage.RequestUri.ToString().Split('/').Last().TrimStart('v');
+                return versionTag;
+            }
+
+            return Optional.None<string>();
+        }
+
+        private Optional<ExecutionPolicies> RetrieveExecutionPolicy()
+        {
+            var process = new Process()
+            {
+                StartInfo = new ProcessStartInfo
+                {
+                    FileName = "powershell",
+                    Arguments = GetExecutionPolicyCommand,
+                    UseShellExecute = false,
+                    RedirectStandardError = true,
+                    RedirectStandardInput = true,
+                    RedirectStandardOutput = true
+                }
+            };
+            process.Start();
+            this.logger.LogInformation("Checking current execution policy");
+            var output = process.StandardOutput.ReadToEnd();
+            if (!Enum.TryParse(typeof(ExecutionPolicies), output, out var executionPolicy))
+            {
+                var error = process.StandardError.ReadToEnd();
+                this.logger.LogError($"Failed to retrieve current user execution policy. Stdout: {output}. Stderr: {error}");
+                return Optional.None<ExecutionPolicies>();
+            }
+
+            return executionPolicy.Cast<ExecutionPolicies>();
+        }
+
+        private void SetExecutionPolicy(ExecutionPolicies executionPolicy)
+        {
+            var process = new Process()
+            {
+                StartInfo = new ProcessStartInfo
+                {
+                    FileName = "powershell",
+                    Arguments = SetExecutionPolicyCommand.Replace(ExecutionPolicyTag, executionPolicy.ToString()),
+                    UseShellExecute = false,
+                    RedirectStandardError = true,
+                    RedirectStandardInput = true,
+                    RedirectStandardOutput = true
+                }
+            };
+            process.Start();
+            this.logger.LogInformation($"Setting execution policy to {executionPolicy}");
+            var output = process.StandardOutput.ReadToEnd();
+            if (!string.IsNullOrWhiteSpace(output))
+            {
+                var error = process.StandardError.ReadToEnd();
+                throw new InvalidOperationException($"Failed to set execution policy to {executionPolicy}. Stdout: {output}. Stderr: {error}");
+            }
+        }
+
+        private void LaunchExtractor()
         {
             File.WriteAllLines(ExtractAndRunPs1, new List<string>()
             {
-                SetExecutionPolicy,
-                DelayCommand,
+                WaitCommand.Replace(ProcessIdTag, Environment.ProcessId.ToString()),
                 ExtractCommandTemplate
                     .Replace(InputFileTag, Path.GetFullPath(TempFile))
                     .Replace(OutputPathTag, Directory.GetCurrentDirectory()),
@@ -124,22 +237,90 @@ public void FinalizeUpdate()
                     WorkingDirectory = Directory.GetCurrentDirectory()
                 },
             };
+            this.logger.LogInformation("Created extractor script. Attempting to launch powershell");
             if (process.Start() is false)
             {
                 throw new InvalidOperationException("Failed to create and start powershell script");
             }
         }
 
-        private async Task<Optional<string>> GetLatestVersion()
+        private static void MarkUpdateInRegistry()
         {
-            using var response = await this.httpClient.GetAsync(Url);
-            if (response.IsSuccessStatusCode)
+            var homeRegistryKey = GetOrCreateHomeKey();
+            homeRegistryKey.SetValue(UpdatedKey, true);
+            homeRegistryKey.Close();
+        }
+
+        private static void UnmarkUpdateInRegistry()
+        {
+            var homeRegistryKey = GetOrCreateHomeKey();
+            homeRegistryKey.SetValue(UpdatedKey, false);
+            homeRegistryKey.Close();
+        }
+
+        private static bool UpdateMarkedInRegistry()
+        {
+            var homeRegistryKey = GetOrCreateHomeKey();
+            var update = homeRegistryKey.GetValue(UpdatedKey);
+            homeRegistryKey.Close();
+            if (update is string updateString)
             {
-                var versionTag = response.RequestMessage.RequestUri.ToString().Split('/').Last().TrimStart('v');
-                return versionTag;
+                if (bool.TryParse(updateString, out var updateValue))
+                {
+                    return updateValue;
+                }
+                else
+                {
+                    throw new InvalidOperationException($"Found update value {updateString} in registry");
+                }
             }
 
-            return Optional.None<string>();
+            return false;
+        }
+
+        private static void SaveExecutionPolicyValueToRegistry(ExecutionPolicies executionPolicy)
+        {
+            var homeRegistryKey = GetOrCreateHomeKey();
+            homeRegistryKey.SetValue(ExecutionPolicyKey, executionPolicy.ToString());
+            homeRegistryKey.Close();
+        }
+
+        private static Optional<ExecutionPolicies> LoadExecutionPolicyValueFromRegistry()
+        {
+            var homeRegistryKey = GetOrCreateHomeKey();
+            var executionPolicy = homeRegistryKey.GetValue(ExecutionPolicyKey);
+            homeRegistryKey.Close();
+            
+            if (executionPolicy is null)
+            {
+                return Optional.None<ExecutionPolicies>();
+            }
+            else if (executionPolicy is string executionPolicyString)
+            {
+                if (Enum.TryParse<ExecutionPolicies>(executionPolicyString, out var executionPolicyValue))
+                {
+                    return executionPolicyValue;
+                }
+                else
+                {
+                    throw new InvalidOperationException($"Found execution policy with value {executionPolicy}");
+                }
+            }
+            else
+            {
+                throw new InvalidOperationException($"Found execution policy of type {executionPolicy.GetType()}.");
+            }
+        }
+
+        private static RegistryKey GetOrCreateHomeKey()
+        {
+            var homeRegistryKey = Registry.CurrentUser.OpenSubKey("Software", true).OpenSubKey(RegistryKey, true);
+            if (homeRegistryKey is null)
+            {
+                homeRegistryKey = Registry.CurrentUser.OpenSubKey("Software", true).CreateSubKey(RegistryKey, true);
+            }
+
+            return homeRegistryKey;
         }
     }
 }

Daybreak/Services/Updater/IApplicationUpdater.cs

@@ -1,9 +1,10 @@
 using Daybreak.Models;
+using Daybreak.Services.ApplicationLifetime;
 using System.Threading.Tasks;
 
 namespace Daybreak.Services.Updater
 {
-    public interface IApplicationUpdater
+    public interface IApplicationUpdater : IApplicationLifetimeService
     {
         string CurrentVersion { get; }
         void FinalizeUpdate();