Public/Restore-AzVmSnapshot.ps1

function Restore-AzVmSnapshot {
    [CmdletBinding(DefaultParameterSetName = 'VM', SupportsShouldProcess, ConfirmImpact = 'High')]
    param
    (
        [Parameter(Mandatory, ParameterSetName = 'VM')]
        [ValidateNotNullOrEmpty()]
        [string]$VmName,
 
        [Parameter(Mandatory, ParameterSetName = 'VM')]
        [ValidateNotNullOrEmpty()]
        [string]$ResourceGroupName,
 
        [Parameter(Mandatory, ValueFromPipeline, ParameterSetName = 'Snapshot')]
        [ValidateNotNullOrEmpty()]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSSnapshotList]$Snapshot,
 
        [Parameter()]
        [ValidateNotNullOrEmpty()]
        [switch]$RemoveOriginalDisk
    )
 
    $ErrorActionPreference = 'Stop'
    
    if (-not $PSBoundParameters.ContainsKey('Snapshot')) {
        ## Find the VM
        $vm = Get-AzVM -Name $VmName -ResourceGroupName $ResourceGroupName
        if (-not ($Snapshot = Get-AzVmSnapshot -ResourceGroupName $ResourceGroupName -VMName $vm.Name)) {
            throw "Could not find snapshot."
        }
    } else {
        $vmName = ($Snapshot.Name -split '-')[1]
        $ResourceGroupName = $Snapshot.ResourceGroupName
        $vm = Get-AzVM -Name $VmName -ResourceGroupName $ResourceGroupName
    }
 
    ## Find the OS disk on the VM to get the storage type
    $osDiskName = $vm.StorageProfile.OsDisk.name
    $oldOsDisk = Get-AzDisk -Name $osDiskName -ResourceGroupName $ResourceGroupName
    
    if ($PSCmdlet.ShouldProcess("Snapshot", 'Restore')) {
        $diskconf = New-AzDiskConfig -AccountType $oldOsDisk.sku.name -Location $oldOsdisk.Location -SourceResourceId $Snapshot.Id -CreateOption Copy
        Write-Verbose -Message 'Creating new disk...'
        $newDisk = New-AzDisk -Disk $diskconf -ResourceGroupName $ResourceGroupName -DiskName "$($vm.Name)-$((New-Guid).ToString())"

        # Set the VM configuration to point to the new disk
        $null = Set-AzVMOSDisk -VM $vm -ManagedDiskId $newDisk.Id -Name $newDisk.Name

        # Update the VM with the new OS disk
        Write-Verbose -Message 'Updating VM...'
        $null = Update-AzVM -ResourceGroupName $ResourceGroupName -VM $vm 

        # Start the VM
        Write-Verbose -Message 'Starting VM...'
        $null = Start-AzVM -Name $vm.Name -ResourceGroupName $ResourceGroupName

        if ($RemoveOriginalDisk.IsPresent) {
            if ($PSCmdlet.ShouldProcess("Disk $($oldOsDisk.Name)", 'Remove')) {
                $null = Remove-AzDisk -ResourceGroupName $ResourceGroupName -DiskName $oldOsDisk.Name
            }
        }
    }
}