Azure Automation Update Management - Execute Pre/Post script on Hybrid Worker Groups with local machines

Execute-RunbookLocally.ps1

<#PSScriptInfo
.VERSION 1.0
.GUID c2984973-98d8-4913-adba-55f1aec1c90f
.AUTHOR justincjahn
.COMPANYNAME Jahn Digital
.COPYRIGHT MIT
.TAGS UpdateManagement, Automation
.LICENSEURI
.PROJECTURI
.ICONURI
.EXTERNALMODULEDEPENDENCIES
.REQUIREDSCRIPTS
.EXTERNALSCRIPTDEPENDENCIES
.RELEASENOTES
.PRIVATEDATA
.NOTES
	Credit to https://github.com/azureautomation/update-management-run-script-locally for
	the original script.
#>

<#
.SYNOPSIS
	Runs a child Automation Runbook on one or more hybrid workers.

.DESCRIPTION
	This script is intended to be run as a part of Update Management Pre/Post scripts.
	It requires hybrid workers to be configured on the machines which need to run scripts locally.
	Runs a child Automation Runbook on a on or more hybrid workers and passes the machines being
	updated to the child Runbook via the TargetMachines parameter.

.PARAMETER RunbookName
	The name of the Azure Automation runbook you wish to execute on the hybrid workers in a local context.

.PARAMETER HybridWorkerGroups
	The hybrid worker group that should run the script.  Can be a comma separated list of Hybrid Worker Groups.

.PARAMETER SoftwareUpdateConfigurationRunContext
	This is a system variable which is automatically passed in by Update Management during a deployment.
#>

param(
	[parameter(Mandatory=$true)]
	[string]
	$RunbookName,

	[parameter(Mandatory=$true)]
	[string]
	$HybridWorkerGroups,

	[string]
	$SoftwareUpdateConfigurationRunContext
)

#region GlobalVariables
	$initialJobId = $PSPrivateMetadata.JobId
	$resourceGroup = $null
	$automationAccount = $null
	$runbookStartFailures = 0
#endregion GlobalVariables

#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

#region GetAutomationAccountInformation
	if ([string]::IsNullOrEmpty($initialJobId)) {
		Write-Output ($PSPrivateMetadata | ConvertTo-Json)
		throw "Unable to obtain the ID of this job.  Exiting."
	}

	try {
		$aAutomationResources = Get-AzResource -ResourceType "Microsoft.Automation/AutomationAccounts"
		foreach ($automationResource in $aAutomationResources) {
			$job = Get-AzAutomationJob `
				-ResourceGroupName $automationResource.ResourceGroupName `
				-AutomationAccountName $automationResource.Name `
				-Id $initialJobId `
				-ErrorAction SilentlyContinue

			if ($null -ne $job) {
				$resourceGroup = $job.ResourceGroupName
				$automationAccount = $job.AutomationAccountName
				break
			}
		}
	} catch {
		Write-Error "Unable to find the automation resource running this job.  Exiting."
		throw $_
	}

	if ($null -eq $resourceGroup -or $null -eq $automationAccount) {
		throw "Unable to find the automation resource running this job.  Exiting."
	}
#endregion GetAutomationAccountInformation

#region ValidateRunbookExists
	try {
		Get-AzAutomationRunbook `
			-ResourceGroupName $resourceGroup `
			-AutomationAccountName $automationAccount `
			-Name $RunbookName `
			-ErrorAction Stop `
		| Out-Null
	} catch {
		throw "The provided runbook, $RunbookName does not exist in the Automation Account.  Exiting."
	}
#endregion ValidateRunbookExists

#region SoftwareUpdateConfigurationContext
	if ([string]::IsNullOrEmpty($SoftwareUpdateConfigurationRunContext)) {
		throw "No Software Update Configuration Run Context provided to the job.  Exiting."
	}

	$context = ConvertFrom-Json $SoftwareUpdateConfigurationRunContext
#endregion SoftwareUpdateConfigurationContext

#region StartJobs
	$machines = $context.SoftwareUpdateConfigurationSettings.NonAzureComputerNames | Select-Object -Unique
	$machines = $machines | ConvertTo-Json
	$aJobs = foreach ($workerGroup in $HybridWorkerGroups.Split(',')) {
		try {
			$job = Start-AzAutomationRunbook `
				-ResourceGroupName $resourceGroup `
				-AutomationAccountName $automationAccount `
				-Name $RunbookName `
				-RunOn $workerGroup `
				-Parameters @{TargetMachines=$machines}

			if ($null -eq $job) {
				[PSCustomObject]@{
					WorkerGroup = $workerGroup
					Job = $null
					Status = "FAILURE"
					Message = "Job did not start.  Unknown error."
				}

				continue
			}

			[PSCustomObject]@{
				WorkerGroup = $workerGroup
				Job = $job
				Status = "RUNNING"
				Message = $null
			}
		} catch {
			[PSCustomObject]@{
				WorkerGroup = $workerGroup
				Job = $null
				Status = "FAILURE"
				Message = $_
			}

			Write-Output $_
		}
	}

	# Loop through and report all jobs that failed
	$aJobs | Where-Object { $_.Status -eq "FAILURE" } | ForEach-Object {
		Write-Error "Failed to start runbook on $($_.WorkerGroup): $($_.Message)."
		$runbookStartFailures++
	}
#endregion StartJobs

#region AwaitJobs
	$aStatuses = @()
	foreach ($job in ($aJobs | Where-Object { $_.Status -ne "FAILURE" })) {
		$currentStatus = $null
		try {
			$currentStatus = Get-AzAutomationJob `
				-ResourceGroupName $resourceGroup `
				-AutomationAccountName $automationAccount `
				-Id $job.Job.JobId
		} catch {
			$aStatuses += [PSCustomObject]@{
				JobId = $job.Job.JobId
				Status = $null
				Output = $null
				Message = "Unable to retrieve automation job: $_"
			}

			continue
		}

		$maxWaitTime = 1800 # 30 Min
		while ($currentStatus.Status -ne "Completed") {
			Write-Output "Waiting for Job $($job.Job.JobId) to complete..."

			if ($maxWaitTime -le 0) {
				Write-Warning "Job $($job.Job.JobId) timed out while waiting for it to complete.  Continuing."
				continue
			}

			Start-Sleep 10
			$maxWaitTime -= 10

			try {
				$currentStatus = Get-AzAutomationJob `
					-ResourceGroupName $resourceGroup `
					-AutomationAccountName $automationAccount `
					-Id $job.Job.JobId
			} catch {
				# We got it once before, try for the full time
				continue
			}
		}

		if ($currentStatus.Status -ne "Completed") {
			Write-Warning "Job $($job.Job.JobId) timed out..."

			$aStatuses += [PSCustomObject]@{
				JobId = $job.Job.JobId
				Status = $currentStatus
				Output = $null
				Message = "Timeout while waiting for completion."
			}

			continue
		}

		$jobOutput = $null
		try {
			$jobOutput = Get-AzAutomationJobOutput `
				-ResourceGroupName $resourceGroup `
				-AutomationAccountName $automationAccount `
				-Id $job.Job.JobId
		} catch {
			$aStatuses += [PSCustomObject]@{
				JobId = $job.Job.JobId
				Status = $currentStatus
				Output = $null
				Message = "Unable to retrieve job output: $_"
			}
		}

		Write-Output "Job $($job.Job.JobId) complete..."

		$aStatuses += [PSCustomObject]@{
			JobId = $job.Job.JobId
			Status = $currentStatus
			Output = $jobOutput
			Message = $null
		}
	}
#endregion AwaitJobs

#region ErrorChecking
	$aErrors = foreach ($status in $aStatuses) {
		if ($null -eq $status.Output) {
			if ($null -eq $status.Message) {
				Write-Output "Job $($status.JobId): An unknown error occurred."
			} else {
				Write-Output "Job $($status.JobId): $($status.Message)."
			}

			continue
		}

		if ($status.Output.Type -eq "Error") {
			Write-Output "Job $($status.JobId): $($status.Output.Summary)."
		}
	}

	if ($aErrors.Length -gt 0 -or $runbookStartFailures -gt 0) {
		$aErrors | ForEach-Object { Write-Error $_ }
		throw "One or more errors ocurred during execution."
	}
#endregion ErrorChecking

Write-Output "Run completed without errors."