From 495f740eda37b8302e06e13d0ea6f7d56d756035 Mon Sep 17 00:00:00 2001
From: Shamil Mubarakshin
 <127750046+shamil-mubarakshin@users.noreply.github.com>
Date: Fri, 10 Nov 2023 16:25:40 +0100
Subject: [PATCH] Adjust services handling in Finalize-VM and helpers (#8771)

---
 .../scripts/ImageHelpers/InstallHelpers.ps1   | 65 ++++++++-----------
 images/win/scripts/Installers/Finalize-VM.ps1 |  6 +-
 2 files changed, 29 insertions(+), 42 deletions(-)

diff --git a/images/win/scripts/ImageHelpers/InstallHelpers.ps1 b/images/win/scripts/ImageHelpers/InstallHelpers.ps1
index a4020253..ad466aaa 100644
--- a/images/win/scripts/ImageHelpers/InstallHelpers.ps1
+++ b/images/win/scripts/ImageHelpers/InstallHelpers.ps1
@@ -92,8 +92,7 @@ function Install-Binary
     }
 }
 
-function Stop-SvcWithErrHandling
-{
+function Stop-SvcWithErrHandling {
     <#
     .DESCRIPTION
         Function for stopping the Windows Service with error handling
@@ -104,36 +103,26 @@ function Stop-SvcWithErrHandling
     .PARAMETER StopOnError
         Switch for stopping the script and exit from PowerShell if one service is absent
     #>
-    Param
-    (
+    Param (
         [Parameter(Mandatory, ValueFromPipeLine = $true)]
         [string] $ServiceName,
         [switch] $StopOnError
     )
 
-    Process
-    {
+    Process {
         $service = Get-Service $ServiceName -ErrorAction SilentlyContinue
-        if (-not $service)
-        {
+        if (-not $service) {
             Write-Warning "[!] Service [$ServiceName] is not found"
-            if ($StopOnError)
-            {
+            if ($StopOnError) {
                 exit 1
             }
-
-        }
-        else
-        {
+        } else {
             Write-Host "Try to stop service [$ServiceName]"
-            try
-            {
+            try {
                 Stop-Service -Name $ServiceName -Force
                 $service.WaitForStatus("Stopped", "00:01:00")
                 Write-Host "Service [$ServiceName] has been stopped successfuly"
-            }
-            catch
-            {
+            } catch {
                 Write-Error "[!] Failed to stop service [$ServiceName] with error:"
                 $_ | Out-String | Write-Error
             }
@@ -141,8 +130,7 @@ function Stop-SvcWithErrHandling
     }
 }
 
-function Set-SvcWithErrHandling
-{
+function Set-SvcWithErrHandling {
     <#
     .DESCRIPTION
         Function for setting the Windows Service parameter with error handling
@@ -152,32 +140,33 @@ function Set-SvcWithErrHandling
 
     .PARAMETER Arguments
         Hashtable for service arguments
+
+    .PARAMETER StopOnError
+        Switch for stopping the script and exit from PowerShell if one service is absent
     #>
 
-    Param
-    (
+    Param (
         [Parameter(Mandatory, ValueFromPipeLine = $true)]
         [string] $ServiceName,
         [Parameter(Mandatory)]
-        [hashtable] $Arguments
+        [hashtable] $Arguments,
+        [switch] $StopOnError
     )
 
-    Process
-    {
+    Process {
         $service = Get-Service $ServiceName -ErrorAction SilentlyContinue
-        if (-not $service)
-            {
-                Write-Warning "[!] Service [$ServiceName] is not found"
+        if (-not $service) {
+            Write-Warning "[!] Service [$ServiceName] is not found"
+            if ($StopOnError) {
+                exit 1
+            }
+        } else {
+            try {
+                Set-Service $serviceName @Arguments
+            } catch {
+                Write-Error "[!] Failed to set service [$ServiceName] arguments with error:"
+                $_ | Out-String | Write-Error
             }
-
-        try
-        {
-           Set-Service $serviceName @Arguments
-        }
-        catch
-        {
-            Write-Error "[!] Failed to set service [$ServiceName] arguments with error:"
-            $_ | Out-String | Write-Error
         }
     }
 }
diff --git a/images/win/scripts/Installers/Finalize-VM.ps1 b/images/win/scripts/Installers/Finalize-VM.ps1
index 398b59c8..8e64259e 100644
--- a/images/win/scripts/Installers/Finalize-VM.ps1
+++ b/images/win/scripts/Installers/Finalize-VM.ps1
@@ -104,10 +104,8 @@ $servicesToDisable = @(
     'StorSvc'
 )
 
-$servicesToDisable | ForEach-Object {
-    Stop-Service -Name $_
-    Set-Service -Name $_ -StartupType Disabled -ErrorAction Ignore
-} | Out-Null
+$servicesToDisable | Stop-SvcWithErrHandling
+$servicesToDisable | Set-SvcWithErrHandling -Arguments @{StartupType = "Disabled"}
 
 # Disable scheduled tasks
 $allTasksInTaskPath = @(