azure-vm-snapshots.ps1

function Invoke-AzVMDiskSnapshot {
    param(
        [Parameter(Mandatory = $true)]
        [string]$VMResourceGroupName,
        [Parameter(Mandatory = $true)]
        [string]$VMName,
        [string]$SnapshotsResourceGroupName = $null
        [bool]$Shutdown = $true
    )

    process {
        if(-not $SnapshotsResourceGroupName -or $SnapshotsResourceGroupName -eq "") {
            $SnapshotsResourceGroupName = $VMResourceGroupName
        }

        $vm = Get-AzVM -ResourceGroupName $VMResourceGroupName -Name $VMName -ErrorAction SilentlyContinue
        if(-not $vm) {
            throw "No VM found"
        }

        Write-Output $vm.Name

        if($Shutdown) {
            Write-Output " - stopping VM"
            Stop-AzVM -ResourceGroupName $vm.ResourceGroupName -Name $vm.Name -Force
        }

        $snapshots = @{}

        # OS disk
        Write-Output " * OS disk: $($vm.StorageProfile.OsDisk.Name)"

        $name = "$($vm.StorageProfile.OsDisk.Name)_snapshot_$(Get-Date -Format ddMMyyHHmm)"
        $snapshots[$name] = New-AzSnapshotConfig -SourceUri $vm.StorageProfile.OsDisk.ManagedDisk.Id -Location $vm.Location -CreateOption Copy

        # Data Disks
        foreach($dataDiskInfo in $vm.StorageProfile.DataDisks) {
            $dataDisk = Get-AzDisk -ResourceGroupName $vm.ResourceGroupName -DiskName $dataDiskInfo.Name

            Write-Output " * Data disk: $($dataDiskInfo.Name)"

            $name = "$($dataDisk.Name)_snapshot_$(Get-Date -Format ddMMyyHHmm)"
            $snapshots[$name] = New-AzSnapshotConfig -SourceUri $dataDisk.Id -CreateOption Copy -Location $vm.Location
        }
        
        foreach($snapshotName in $snapshots.Keys) {
            New-AzSnapshot -Snapshot $snapshots[$snapshotName] -SnapshotName $snapshotName -ResourceGroupName $SnapshotsResourceGroupName 
        }
    }
}

$VMResourceGroupName = "hcittests"
$VMName = "datavm"
$SnapshotsResourceGroupName = "hcitests-snapshots"

Invoke-AzVMDiskSnapshot -VMResourceGroupName $VMResourceGroupName -VMName $VMName -SnapshotsResourceGroupName $SnapshotsResourceGroupName

function Restore-AzVMFromDiskSnapshot {
    param(
        [Parameter(Mandatory = $true)]
        [string]$VMResourceGroupName,
        [Parameter(Mandatory = $true)]
        [string]$VMName,
        [Parameter(Mandatory = $true)]
        [string]$OsDiskSnapshotName,
        [string]$DataDisksSnapshotNames,
        [string]$SnapshotsResourceGroupName = $null
    )

    process {
        $vm = Get-AzVM -ResourceGroupName $VMResourceGroupName -Name $VMName -ErrorAction SilentlyContinue
        if(-not $vm) {
            throw "No VM found"
        }

        Write-Output $vm.Name

        $dataDiskSnapshots = $dataDisksSnapshotNames -split ","
        if($dataDiskSnapshots.Count -ne $vm.StorageProfile.DataDisks.Count) {
            throw "Data disk snapshots does not correspond to number of currenctly attached data disks!"
        }

        # Check input parameters first
        $osDiskSnapshotInfo = Get-AzSnapshot -ResourceGroupName $SnapshotsResourceGroupName -SnapshotName $OsDiskSnapshotName
        if(-not $osDiskSnapshotInfo) {
            throw "Cannot find OS Disk snapshot $($OsDiskSnapshotName)!"
        }

        $dataDisksInfo = @()
        foreach($dataDisksSnapshotName in $dataDiskSnapshots) {
            $dataDiskSnapshotInfo = Get-AzSnapshot -ResourceGroupName $SnapshotsResourceGroupName -SnapshotName $dataDisksSnapshotName
            if(-not $dataDiskSnapshotInfo) {
                throw "Cannot find Data disk snapshot $($dataDisksSnapshotName)!"
            }

            $dataDisksInfo += $dataDiskSnapshotInfo
        }

        Write-Output " - creating disks from snapshot"

        # Create OS disk
        $osDiskName = $OsDiskSnapshotName -replace "_snapshot_", ""
        $test = Get-AzDisk -ResourceGroupName $vm.ResourceGroupName -DiskName $osDiskName -ErrorAction SilentlyContinue
        if($test) {
            throw "This snapshot seem to be already applied!"
        }

        $osDisk = New-AzDisk -DiskName $osDiskName (New-AzDiskConfig -Location $vm.Location -CreateOption Copy -SourceResourceId $osDiskSnapshotInfo.Id) -ResourceGroupName $vm.ResourceGroupName

        # Create data disks
        $dataDisks = @()
        foreach($dataDiskSnapshotInfo in $dataDisksInfo) {
            $diskName = $dataDiskSnapshotInfo.Name -split "_snapshot_" | Select-Object -First 1
            $currentDisk = $vm.StorageProfile.DataDisks | Where-Object Name -EQ $diskName
            if(-not $currentDisk) {
                throw "Cannot find current disk $($diskName) on VM to detect Lun order"
            }

            $dataDiskName = $dataDiskSnapshotInfo.Name -replace "_snapshot_", ""
            $dataDisk = New-AzDisk -DiskName $dataDiskName (New-AzDiskConfig -Location $vm.Location -CreateOption Copy -SourceResourceId $dataDiskSnapshotInfo.Id) -ResourceGroupName $vm.ResourceGroupName

            $dataDisks += [PSCustomObject]@{
                ParentDiskName = $currentDisk.Name
                Lun = $currentDisk.Lun
                Name = $dataDiskName
                Id = $dataDisk.Id
            }
        }

        # Stop the VM
        Write-Output " - stopping VM"
        Stop-AzVM -ResourceGroupName $vm.ResourceGroupName -Name $vm.Name -Force | Out-Null

        Write-Output " - attaching new disks"

        # Swap OS Disk
        Set-AzVMOSDisk -VM $vm -ManagedDiskId $osDisk.Id -Name $osDisk.Name

        Update-AzVM -ResourceGroupName $vm.ResourceGroupName -VM $vm

        # replace data disks
        foreach($dataDisk in $dataDisks) {
            Write-Output " - replacing $($dataDisk.ParentDiskName) with $($dataDisk.Name)"
            Remove-AzVMDataDisk -VM $vm -DataDiskNames $dataDisk.ParentDiskName
            # replace is not enabled, so we need to commit removal first, then we can attach again to same lun
            Update-AzVM -ResourceGroupName $vm.ResourceGroupName -VM $vm
            
            Add-AzVMDataDisk -VM $vm -Name $dataDisk.Name -CreateOption Attach -ManagedDiskId $dataDisk.Id -Lun $dataDisk.Lun
            Update-AzVM -ResourceGroupName $vm.ResourceGroupName -VM $vm
        }

        # And start VM back
        Write-Output " - starting VM"
        Start-AzVM -Name $vm.Name -ResourceGroupName $vm.ResourceGroupName
    }
}

$VMResourceGroupName = "hcittests"
$VMName = "datavm"
$SnapshotsResourceGroupName = "hcitests-snapshots"
$OsDiskSnapshotName = "datavm_disk1_1e2e6f344b1a4cbc8268709a12bebf37_snapshot_2907201022"
$dataDisksSnapshotNames = "datavm_DataDisk_0_snapshot_2907201022"

Restore-AzVMFromDiskSnapshot -VMResourceGroupName $VMResourceGroupName -VMName $VMName -SnapshotsResourceGroupName $SnapshotsResourceGroupName -OsDiskSnapshotName $OsDiskSnapshotName -dataDisksSnapshotNames $dataDisksSnapshotNames