AzureUploadClient.ps1


Class AzureUploadClient : UploadClient {
    # properties
    [string]$AzureSubscriptionId
    [string]$AzureLocation
    [string]$TargetResourceGroup
    [string]$CloudDiskName
    [int]$UploadTimeout

    [string]$AzureStorageType = 'Premium_LRS'
    [int]$Threads = 5


    # Default constructor
    AzureUploadClient([hashtable]$values) : base ($values) {
        $this.CopyValues(@{
            AzureSubscriptionId = @{
                required = $True
            }
            AzureLocation = @{
                required = $True
            }
            TargetResourceGroup = @{
                required = $True
            }
            UploadTimeout = @{
                default = 36000 
                required = $False
            }
        }, $values)
    }

    Cleanup() {
        LogIt "Deleting existing disk $($this.CloudDiskName) from Azure subscription $($this.AzureSubscriptionId) resource group $($this.TargetResourceGroup)"
        SelectAzSubscription $this.AzureSubscriptionId

        LogIt "Finding existing disks named $($this.CloudDiskName)"
        Get-AzDisk -ResourceGroupName $this.TargetResourceGroup -DiskName "$($this.CloudDiskName)*" | ForEach-Object {
            $diskName = $_.Name
            LogIt "Deleting existing managed disk $diskName"
            $result = Remove-AzDisk -ResourceGroupName $this.TargetResourceGroup -DiskName $diskName -Force
            if ($null -eq $result) {
                LogFatal "Failed to delete existing managed disk $diskName"
            }
            LogIt "Deleted existing managed disk $diskName"
        }
    }

    hidden [string]CreateManagedDisk([long]$sizeInBytes)
    {
        $err = ""

        $SasExpiryDuration = $this.UploadTimeout
 
        SelectAzSubscription $this.AzureSubscriptionId

        $diskConfig = New-AzDiskConfig -AccountType $this.AzureStorageType -Location $this.AzureLocation -UploadSizeInBytes $sizeInBytes -CreateOption 'Upload' -OsType Windows -HyperVGeneration V2

        LogIt "Creating managed disk $($this.CloudDiskName) in sub $($this.AzureSubscriptionId) rg $($this.TargetResourceGroup) loc $($this.AzureLocation) with size $sizeInBytes bytes"
        $result = New-AzDisk -ResourceGroupName $this.TargetResourceGroup -DiskName $this.CloudDiskName -Disk $diskConfig -ErrorVariable "err"
        if ($null -eq $result)
        {
            throw "Failed to create managed disk $($this.CloudDiskName) in rg $($this.TargetResourceGroup) error: $err"
        }

        LogIt "Granting access to managed disk $($this.CloudDiskName) for $($SasExpiryDuration) seconds"
        $access = Grant-AzDiskAccess -ResourceGroupName $this.TargetResourceGroup -DiskName $this.CloudDiskName -DurationInSecond $SasExpiryDuration -Access 'Write' -ErrorVariable "err"
        if ($null -eq $access)
        {
            throw "Failed to create sas for maanged disk $($this.CloudDiskName) error: $err"
        }

        $sas = $access.AccessSAS
        LogIt "Created managed disk $($this.CloudDiskName) with sas $sas"
        return $sas
    }

    hidden [void]CopyWithCloudUploader([string]$destination, [string]$smbPassword) {
        $cloudupload = {
            param($SharePath, $Arguments)
            $fullPath = Join-Path -Path $SharePath -ChildPath $Arguments[0]
            LogIt "Copying $fullPath to $destination with CloudUploader using $($this.Threads) threads"
            Copy-ToAzDisk -File $fullPath -Sas $destination -Threads $this.Threads
        }
        ExecuteOnSmbShareAsUser $cloudupload $this.smbConfig.UserAndDomain $smbPassword $this.smbConfig.ShareUnc @($this.smbConfig.ExportFilePath)
    }

    hidden [void]CopyWithAzcopy([string]$destination, [string]$smbPassword) {
        $azcopy = {
            param($SharePath, $Arguments)
            $fullPath = Join-Path -Path $SharePath -ChildPath $Arguments[0]
            $log = Get-Location
            LogIt "Copying $fullPath to $destination with azcopy, logging to $log"
            $env:AZCOPY_LOG_LOCATION = $log
            $copy = azcopy copy $fullPath $destination --blob-type PageBlob
            LogIt "Result from azcopy: $copy"
            if ($LastExitCode -ne 0) {
                throw "azcopy of $($this.CloudDiskName) failed ($LastExitCode)."
            }
        }
        ExecuteOnMappedShareAsUser $azcopy $this.smbConfig.UserAndDomain $smbPassword $this.smbConfig.ShareUnc @($this.smbConfig.ExportFilePath)
    }

    # Upload exported image from a share to Azure.
    [psobject]Upload([string]$smbPassword, [bool]$azcopy) {
        if ([string]::IsNullOrWhiteSpace($smbPassword)) {
            throw "Upload requires SMB share user password to be supplied."
        }

        LogIt "Uploading disk to Azure"

        $disk = $null
    
        $fileSize = $this.GetVhdSize($smbPassword)
        $managedDiskUrlWithSas = $this.CreateManagedDisk($fileSize)

        try
        {
            if ($azcopy) {
                $this.CopyWithAzcopy($managedDiskUrlWithSas, $smbPassword)
            } else {
                $this.CopyWithCloudUploader($managedDiskUrlWithSas, $smbPassword)
            }
            $disk = Get-AzDisk -ResourceGroupName $this.TargetResourceGroup -DiskName $this.CloudDiskName
            LogIt "Created and uploaded disk $disk"
        } catch {
            Remove-AzDisk -ResourceGroupName $this.TargetResourceGroup -DiskName $this.CloudDiskName -Force
            LogFatal "Disk copy failed: $_"
        }
        finally
        {
            $revoke = Revoke-AzDiskAccess -ResourceGroupName $this.TargetResourceGroup -DiskName $this.CloudDiskName
            LogIt "Revoke-AzDiskAccess status: $($revoke.status)"
        }

        return @{
            diskName = $disk.Name
            resourceGroup = $this.TargetResourceGroup
        }
    }
}