<#
.SYNOPSIS
 Barebones script for Update Management Pre/Post

.DESCRIPTION
  This script is intended to be run as a part of Update Management pre/post-scripts.
  It requires the Automation account's system-assigned managed identity.

.PARAMETER SoftwareUpdateConfigurationRunContext
  This is a system variable which is automatically passed in by Update Management during a deployment.
#>

param(
    [string]$SoftwareUpdateConfigurationRunContext
)

#region BoilerplateAuthentication
# Ensures you do not inherit an AzContext in your runbook
Disable-AzContextAutosave -Scope Process

# Connect to Azure with system-assigned managed identity
$AzureContext = (Connect-AzAccount -Identity).context

# set and store context
$AzureContext = Set-AzContext -SubscriptionName $AzureContext.Subscription -DefaultProfile $AzureContext
#endregion BoilerplateAuthentication

#If you wish to use the run context, it must be converted from JSON
$context = ConvertFrom-Json $SoftwareUpdateConfigurationRunContext
#Access the properties of the SoftwareUpdateConfigurationRunContext
$vmIds = $context.SoftwareUpdateConfigurationSettings.AzureVirtualMachines | Sort-Object -Unique
$runId = $context.SoftwareUpdateConfigurationRunId

Write-Output $context

#----------------------------------------
#The snippet below is based on the sample available at
# https://github.com/azureautomation/update-management-run-script-with-run-command/blob/master/UpdateManagement-RunCommand.ps1
#----------------------------------------

#If you wish to use the run context, it must be converted from JSON
$context = ConvertFrom-Json  $SoftwareUpdateConfigurationRunContext
$vmIds = $context.SoftwareUpdateConfigurationSettings.AzureVirtualMachines
$runId = $context.SoftwareUpdateConfigurationRunId

if (!$vmIds) 
{
    #Workaround: Had to change JSON formatting
    $Settings = ConvertFrom-Json $context.SoftwareUpdateConfigurationSettings
    #Write-Output "List of settings: $Settings"
    $VmIds = $Settings.AzureVirtualMachines
    #Write-Output "Azure VMs: $VmIds"
    if (!$vmIds) 
    {
        Write-Output "No Azure VMs found"
        return
    }
}

#----------------------------------------
#The script you wish to run on each VM
$scriptBlock = @'
try{
    $service_name = 'service1'  #update first service name here.
    start-service -name $service_name  #starting service 1
    $temp_ctr = 0
    $success = $false

    while($temp_ctr -lt 3) #wait for 100 milliseconds before trying again
    {
        $service_status = get-service -name $service_name
        if($service_status -eq 'Running')
        {
            $success = $true
            break
        }
        Start-Sleep -Milliseconds 100 #sleep for 100 milliseconds
        $temp_ctr++ #increase cout
    }
    if(!$success)  #if starting the service1 failed,
    {
        throw "$service_name did not start"
    }

    #-------------------------------------------------
    #else, proceed to start serice 2
    $service_name = 'service2'  #update second service name here.
    start-service -name $service_name  #starting service 2
    $temp_ctr = 0
    $success = $false

    while($temp_ctr -lt 3) #wait for 100 milliseconds before trying again
    {
        $service_status = get-service -name $service_name
        if($service_status -eq 'Running')
        {
            $success = $true
            break
        }
        Start-Sleep -Milliseconds 100 #sleep for 100 milliseconds
        $temp_ctr++ #increase cout
    }
    if(!$success)  #if starting the service1 failed,
    {
        throw "$service_name did not start"
    }

    #-------------------------------------------------
    #else, proceed to start serice 3, 
    #...
    #...



}
catch
{
    #throw the exception caught, if any
    throw $_
}

'@
#----------------------------------------

$scriptPath = "$runID.ps1"
#The cmdlet only accepts a file, so temporarily write the script to disk using runID as a unique name
Out-File -FilePath $scriptPath -InputObject $scriptBlock
$scriptFile = get-item $scriptpath
$fullPath = $scriptfile.fullname

$jobIDs= New-Object System.Collections.Generic.List[System.Object]

#Start script on each machine
$vmIds | ForEach-Object {
    $vmId =  $_
    
    $split = $vmId -split "/";
    $subscriptionId = $split[2]; 
    $rg = $split[4];
    $name = $split[8];
    Write-Output ("Subscription Id: " + $subscriptionId)
    $mute = Select-AzureRmSubscription -Subscription $subscriptionId
    Write-Output "Invoking command on '$($name)' ..."
    $newJob = Start-ThreadJob -ScriptBlock { param($resourceGroup, $vmName, $scriptPath) Invoke-AzureRmVMRunCommand -ResourceGroupName $resourceGroup -Name $VmName -CommandId 'RunPowerShellScript' -ScriptPath $scriptPath} -ArgumentList $rg, $name, $fullPath
    $jobIDs.Add($newJob.Id)
    
}

$jobsList = $jobIDs.ToArray()
if ($jobsList)
{
    Write-Output "Waiting for machines to finish executing..."
    Wait-Job -Id $jobsList
}

foreach($id in $jobsList)
{
    $job = Get-Job -Id $id
    if ($job.Error)
    {
        Write-Output $job.Error
    }

}

#Clean up our variables:
Remove-Item -Path "$runID.ps1"