Copy-Disk.ps1

<#
 .Synopsis
  Implement a cloud disk upload job

 .Description
  Implement a cloud disk upload job. This function supports copy disk from smb share
  server to multiple platform (Azure, AWS and GCP)
#>

Function Copy-Disk
{
    [CmdletBinding(DefaultParameterSetName = 'cmd')]
    Param(
        [Parameter(Mandatory = $true, ParameterSetName = 'file')]
        [string] $ConfigJsonFile,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $CustomerId,
        [Parameter(Mandatory = $true, ParameterSetName = 'cmd')]
        [string] $CloudPlatform,
        [Parameter(Mandatory = $true, ParameterSetName = 'cmd')]
        [string] $SmbHost,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $SmbPort = $null,
        [Parameter(Mandatory = $true, ParameterSetName = 'cmd')]
        [string] $SmbShare,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $SmbPath,
        [Parameter(Mandatory = $true, ParameterSetName = 'cmd')]
        [string] $SmbDiskName,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $SmbDiskFormat = "VhdDiskFormat",
        [Parameter(Mandatory = $true, ParameterSetName = 'cmd')]
        [string] $SmbUserDomain,
        [Parameter(Mandatory = $true, ParameterSetName = 'cmd')]
        [string] $SmbUserName,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $AzureSubscriptionId,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $AzureLocation,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $TargetResourceGroup,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $CloudDiskName,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [int] $UploadTimeout = 36000,
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [string] $AzureStorageType = "Premium_LRS",
        [Parameter(Mandatory = $false, ParameterSetName = 'cmd')]
        [int] $Threads = 5,
        [Parameter(Mandatory = $false)]
        [string] $AzureClientId,
        [Parameter(Mandatory = $false)]
        [string] $AzureSecret,
        [Parameter(Mandatory = $false)]
        [string] $AzureTenantId,
        [Parameter(Mandatory = $false)]
        [string] $AwsRegion,
        [Parameter(Mandatory = $false)]
        [string] $AwsProfileName,
        [Parameter(Mandatory = $false)]
        [string] $GcpServiceAccountKeyFile,
        [Parameter(Mandatory = $false)]
        [pscredential] $SmbCred,
        [Parameter(Mandatory = $false)]
        [switch] $Install,
        [Parameter(Mandatory = $false)]
        [switch] $OverwriteLog,
        [Parameter(Mandatory = $false)]
        [switch] $Force
    )

    Begin
    {
        InitUploadLog $OverwriteLog
        if ($PSCmdlet.ParameterSetName -eq 'file') {
            Log "Loading config from $ConfigJsonFile"
            $configData = Get-Content -Raw -Path $ConfigJsonFile | ConvertFrom-Json
            Log "Config: $configData" $False
            $CustomerId = $configData.CustomerId
            $CloudPlatform = $configData.CloudPlatform
            $SmbHost = $configData.UploadSmb.Host
            $SmbPort = $configData.UploadSmb.Port
            $SmbShare = $configData.UploadSmb.Share
            $SmbPath = $configData.UploadSmb.Path
            $SmbDiskName = $configData.UploadSmb.DiskName
            $SmbDiskFormat = $configData.UploadSmb.DiskFormat
            if([String]::IsNullOrWhiteSpace($SmbDiskFormat)) { $SmbDiskFormat = "VhdDiskFormat" }
            $SmbUserDomain = $configData.UploadSmb.UserDomain
            $SmbUserName = $configData.UploadSmb.UserName
            $AzureSubscriptionId = $configData.AzureSubscriptionId
            $AzureLocation = $configData.AzureLocation
            $TargetResourceGroup = $configData.TargetResourceGroup
            $CloudDiskName = $configData.CloudDiskName
            if($configData.UploadTimeout -eq $null -Or $configData.UploadTimeout -le 0) { $UploadTimeout = 36000 }
            else { $UploadTimeout = [int]$configData.UploadTimeout }
            if([String]::IsNullOrWhiteSpace($configData.AzureStorageType)) { $AzureStorageType = "Premium_LRS" }
            else {$AzureStorageType = $configData.AzureStorageType}
            if($configData.Threads -eq $null -Or $configData.Threads -le 0) { $Threads = 5 }
            else { $Threads = [int]$configData.Threads }
        }
        Log "***** Initialize Smb Config *****"
        $smbConfig = InitSmbConfig $SmbHost $SmbPort $SmbShare $SmbPath $SmbUserDomain $SmbUserName $SmbDiskName $SmbDiskFormat
        if ($SmbCred -eq $null) {
            Log "Generating Smb Credential using username and password"
            $password = Read-Host -assecurestring "SmbUserPassword"
            $smbConfig.SmbCred = New-Object -TypeName System.Management.Automation.PSCredential -ArgumentList $smbConfig.UserAndDomain, $password
        } else {
            Log "Smb Credential is given as input"
            $smbConfig.SmbCred = $SmbCred
        }
    }

    Process
    {
        if ($CloudPlatform -eq "Azure")
        {
            Log "***** Load Required Modules *****"
            LoadModules @('Az.Accounts', 'Az.Compute') $Install
            Log "***** Authenticate *****"
            AuthAzure $AzureClientId $AzureSecret $AzureTenantId $AzureSubscriptionId
            if ($Force) {
                Log "***** Pre Upload Cleanup *****"
                CleanUpAzureDisk $AzureSubscriptionId $CloudDiskName $TargetResourceGroup
            }
            Log "***** Starting Upload *****"
            Log "1. Get Vhd Size"
            $fileSize = GetVhdSize $smbConfig
            Log "2. Create Managed Disk"
            $managedDiskUrlWithSas = CreateManagedDisk $fileSize $UploadTimeout $AzureSubscriptionId $AzureStorageType $AzureLocation $TargetResourceGroup $CloudDiskName
            Log "3. Uploading"
            $InformationPreference = "Continue"
            UploadFromSmbToAzure $managedDiskUrlWithSas $smbConfig $Threads $TargetResourceGroup $CloudDiskName
            Log "***** Finished Upload *****"
        }

        if ($CloudPlatform -eq "Aws")
        {
            Log "***** Starting Upload *****"
            $InformationPreference = "Continue"
            try
            {
                UploadFromSmbToAws $smbConfig $AwsProfileName $AwsRegion $Threads
            }
            catch
            {
                Log "Exception encountered when uploading"
                Log $_.Exception $false
                ThrowException "Failed to upload disk to AWS : $_"
            }
            Log "***** Finished Upload *****"
        }

        if ($CloudPlatform -eq "Gcp")
        {
            # Temporary solution to redirecting 1.48 assemblies that are compile time referenced in the Google.Api.Gax.Rest assemblies
            # to the 1.49 assemblies that the other Google assemblies reference
            # This is the powershell way to accomplish Binding Redirects, that are typically done in the configuration file.
            $modulePath = (Get-Item (Get-Module -Name Citrix.Image.Uploader).Path).DirectoryName
            $GoogleApis = [reflection.assembly]::LoadFrom($modulePath + "\bin\netstandard2.0\Google.Apis.dll")
            $GoogleApisCore = [reflection.assembly]::LoadFrom($modulePath + "\bin\netstandard2.0\Google.Apis.Core.dll")
            $GoogleApisAuth = [reflection.assembly]::LoadFrom($modulePath + "\bin\netstandard2.0\Google.Apis.Auth.dll")
            $OnAssemblyResolve = [System.ResolveEventHandler] {
                param($s, $e)
                Log "Resolving Assembly '$($e.Name)'" $false
                if (($e.Name.StartsWith("Google.Apis.Core, Version=1.49.0.0")) -or
                    ($e.Name.StartsWith("Google.Apis.Auth, Version=1.49.0.0")) -or
                    ($e.Name.StartsWith("Google.Apis, Version=1.49.0.0"))) {
                    Log "This workaround may no longer be necessary. The Google Assemblies may now be referencing the correct assemblies. Try removing this event handler (OnAssemblyResolve) and try again."  $false
                }
                if ($e.Name -eq "Google.Apis.Core, Version=1.48.0.0, Culture=neutral, PublicKeyToken=4b01fa6e34db77ab") {
                    Log "Forcing the Assembly '$($e.Name)' to be '$($GoogleApisCore.GetName().Name), Version=$($GoogleApisCore.GetName().Version)'" $false
                    return $GoogleApisCore
                }
                if ($e.Name -eq "Google.Apis.Auth, Version=1.48.0.0, Culture=neutral, PublicKeyToken=4b01fa6e34db77ab") {
                    Log "Forcing the Assembly '$($e.Name)' to be '$($GoogleApisAuth.GetName().Name), Version=$($GoogleApisAuth.GetName().Version)'" $false
                    return $GoogleApisAuth
                }
                if ($e.Name -eq "Google.Apis, Version=1.48.0.0, Culture=neutral, PublicKeyToken=4b01fa6e34db77ab") {
                    Log "Forcing the Assembly '$($e.Name)' to be '$($GoogleApis.GetName().Name), Version=$($GoogleApis.GetName().Version)'" $false
                    return $GoogleApis
                }
                foreach($a in [System.AppDomain]::CurrentDomain.GetAssemblies()) {
                    if ($a.FullName -eq $e.Name) {
                        return $a
                    }
                }
                return $null
            }
            Log "Registering AssemblyResolve Event Handling."
            [System.AppDomain]::CurrentDomain.add_AssemblyResolve($OnAssemblyResolve)

            if ($Force) {
                Log "***** Pre Upload Cleanup *****"
                CleanUpGcpDisk $GcpServiceAccountKeyFile $CloudDiskName
            }
            Log "***** Starting Upload *****"
            $InformationPreference = "Continue"
            try {
                UploadFromSmbToGcp $CloudDiskName $GcpServiceAccountKeyFile $smbConfig
            } catch [System.Reflection.ReflectionTypeLoadException] {
                Log "Message: $($_.Exception.Message)" $false
                Log "StackTrace: $($_.Exception.StackTrace)" $false
                Log "LoaderExceptions: $($_.Exception.LoaderExceptions)" $false
                try {
                    Log "Redirected Google SDK package version and retry uploading process"
                    UploadFromSmbToGcp $CloudDiskName $GcpServiceAccountKeyFile $smbConfig
                }
                catch {
                    Log "Exception encountered when uploading"
                    Log $_.Exception $false
                    ThrowException "Failed to upload disk $CloudDiskName to $CloudPlatform"
                }
            } catch {
                Log "Exception encountered when uploading"
                Log $_.Exception $false
                ThrowException "Failed to upload disk $CloudDiskName to $CloudPlatform"
            }
        }
    }
}

Function InitSmbConfig([string]$smbHost, [string]$smbPort, [string]$smbShare, [string]$smbPath, [string]$smbUserDomain, [string]$smbUserName, [string]$smbDiskName, [string]$smbDiskFormat) {
    $DISK_FORMATS = @{
        VhdDiskFormat = "vhd"
        VhdxDiskFormat = "vhdx"
        VmdkDiskFormat = "vmdk"
        VmdkSparseDiskFormat = "vmdk"
        QCow2DiskFormat = "qcow"
        RawDiskFormat = "raw"
    }
    $smbConfig = @{}
    $smbConfig.DiskExtension = $DISK_FORMATS[$smbDiskFormat]
    $smbConfig.UserAndDomain = "$($smbUserDomain)\$($smbUserName)"

    if ($smbPort) {
        $smbConfig.ShareUnc = "\\$($smbHost):$($smbPort)\$($smbShare)"
    } else {
        $smbConfig.ShareUnc = "\\$($smbHost)\$($smbShare)"
    }
    if ($smbPath) {
        $smbConfig.ExportFilePath = Join-Path -Path $smbPath -ChildPath "$($smbDiskName).$($smbConfig.DiskExtension)"
    } else {
        $smbConfig.ExportFilePath = "$($smbDiskName).$($smbConfig.DiskExtension)"
    }
    $smbConfig.FileOnShare = Join-Path -Path $smbConfig.ShareUnc -ChildPath $smbConfig.ExportFilePath
    return $smbConfig
}

Function LoadModules([string[]]$modules, [bool]$installModule = $False) {
    foreach ($module in $modules) {
        if (Get-Module -ListAvailable -Name $module) {
            Log "Module $module exists"
        }
        else {
            if ($installModule) {
                Log "Installing $module"
                Install-Module -Name $module -Scope CurrentUser -AllowClobber -Force
            } else {
                ThrowException "Module $module is missing. You can either install it manually or add -Install in the cmdlet to auto-install all missing required modules"
            }
        }
    }
}

Function AuthAzure([string]$azureClientId, [string]$azureSecret, [string]$azureTenantId, [string]$azureSubscriptionId) {
    if ($azureClientId -And $azureSecret -And $azureTenantId) {
        Log "Authenticating to Azure using ServicePrincipal id $azureClientId"
        $secret = ConvertTo-SecureString $azureSecret -AsPlainText -Force
        $pscredential = New-Object -TypeName System.Management.Automation.PSCredential($azureClientId, $secret)
        $azureLoginResult = Connect-AzAccount -ServicePrincipal -Credential $pscredential -Tenant $azureTenantId
    } else {
        Log "Authenticating to Azure using interactive auth"
        $context = Get-AzContext
        if (!$context -or ($context.Subscription.Id -ne $azureSubscriptionId)) {
            $output = Connect-AzAccount -Subscription $azureSubscriptionId
            if (!$output) {
                Log "Did not authenticate with Azure" $False
                $azureLoginResult = $False
            } else {
                Log "SubscriptionId '$azureSubscriptionId' connected"
                $azureLoginResult = $True
            }
        } else {
            Log "SubscriptionId '$azureSubscriptionId' already connected"
            $azureLoginResult = $True
        }
    }
    if (!$azureLoginResult) {
        Log "Authenticated to Azure failed!"
    }
    Log "Authenticated to Azure ($azureLoginResult)"
}

Function CleanUpAzureDisk([string]$azureSubscriptionId, [string]$cloudDiskName, [string]$targetResourceGroup) {
    $err = ""
    Log "Deleting existing disk $($cloudDiskName) from Azure subscription $($azureSubscriptionId) resource group $($targetResourceGroup)"
    $azContextResult = Select-AzSubscription -SubscriptionId $azureSubscriptionId -ErrorVariable "err"
    if ($null -eq $azContextResult) {
        ThrowException "Failed to find azure subscription $azureSubscriptionId in AzContext error: $err"
    }
    Get-AzDisk -ResourceGroupName $targetResourceGroup -DiskName $cloudDiskName | ForEach-Object {
        $diskName = $_.Name
        Log "Deleting existing managed disk $diskName"
        $result = Remove-AzDisk -ResourceGroupName $targetResourceGroup -DiskName $diskName -Force -ErrorVariable "err"
        if ($null -eq $result) {
            ThrowException "Failed to delete existing managed disk $diskName error: $err"
        }
        Log "Deleted existing managed disk $diskName"
    }
}

Function GetVhdSize([psobject]$smbConfig) {
    $err = ""
    try {
        Log "Getting VHD size as $($smbConfig.UserAndDomain) for $($smbConfig.FileOnShare)"
        $getVhdSize = {
            param($SharePath, $Arguments)
            $fullPath = Join-Path -Path $SharePath -ChildPath $Arguments[0]
            return Get-VhdSize -File $fullPath -RoundUp -IncludeFooterSize -ErrorVariable "err"
        }
        $fileSize = ExecuteOnSmbShare $getVhdSize $smbConfig.SmbCred $smbConfig.ShareUnc @($smbConfig.ExportFilePath)
        Log "VHD size for $($smbConfig.FileOnShare) is $fileSize"
        return $fileSize
    }
    catch {
        ThrowException "Failed to get VHD size for $($smbConfig.FileOnShare) Get-VhdSize error $err : $_"
    }
    return -1
}

Function CreateManagedDisk([long]$sizeInBytes, [int]$uploadTimeout, [string]$azureSubscriptionId,
    [string]$azureStorageType, [string]$azureLocation, [string]$targetResourceGroup, [string]$cloudDiskName) {
    $err = ""
    $sasExpiryDuration = $uploadTimeout
    $azContextResult = Select-AzSubscription -SubscriptionId $azureSubscriptionId -ErrorVariable "err"
    if ($null -eq $azContextResult) {
        ThrowException "Failed to find azure subscription $azureSubscriptionId in AzContext error: $err"
    }
    $diskConfig = New-AzDiskConfig -AccountType $azureStorageType -Location $azureLocation -UploadSizeInBytes $sizeInBytes -CreateOption 'Upload' -OsType Windows -HyperVGeneration V2
    Log "Creating managed disk $($cloudDiskName) in sub $($azureSubscriptionId) rg $($targetResourceGroup) loc $($azureLocation) with size $sizeInBytes bytes"
    $result = New-AzDisk -ResourceGroupName $targetResourceGroup -DiskName $cloudDiskName -Disk $diskConfig -ErrorVariable "err"
    if ($null -eq $result)
    {
        ThrowException "Failed to create managed disk $($cloudDiskName) in rg $($targetResourceGroup) error: $err"
    }

    Log "Granting access to managed disk $($cloudDiskName) for $($sasExpiryDuration) seconds"
    $access = Grant-AzDiskAccess -ResourceGroupName $targetResourceGroup -DiskName $cloudDiskName -DurationInSecond $sasExpiryDuration -Access 'Write' -ErrorVariable "err"
    if ($null -eq $access)
    {
        ThrowException "Failed to create sas for mananged disk $($cloudDiskName) error: $err"
    }
    $sas = $access.AccessSAS
    Log "Created managed disk $($cloudDiskName) with sas $sas"
    return $sas
}

Function UploadFromSmbToAzure([string]$destination, [psobject]$smbConfig, [int]$threads, [string]$targetResourceGroup, [string]$cloudDiskName) {
    try {
        CopyWithCloudUploader $destination $smbConfig $threads
        $disk = Get-AzDisk -ResourceGroupName $targetResourceGroup -DiskName $cloudDiskName
        Log "Created and uploaded disk $($disk.name)"
    } catch {
        Remove-AzDisk -ResourceGroupName $targetResourceGroup -DiskName $cloudDiskName -Force
        ThrowException "Disk copy failed: $_"
    } finally {
        Log "Revoking Azure Disk Access"
        $revoke = Revoke-AzDiskAccess -ResourceGroupName $targetResourceGroup -DiskName $cloudDiskName
        Log "Revoke-AzDiskAccess status: $($revoke.status)"
    }
}

Function CopyWithCloudUploader([string]$destination, [psobject]$smbConfig, [int]$threads) {
    $cloudupload = {
        param($SharePath, $Arguments)
        $fullPath = Join-Path -Path $SharePath -ChildPath $Arguments[0]
        Log "Copying $fullPath to $destination with CloudUploader using $threads threads"
        Copy-ToAzDisk -File $fullPath -Sas $destination -Threads $threads
    }
    ExecuteOnSmbShare $cloudupload $smbConfig.SmbCred $smbConfig.ShareUnc @($smbConfig.ExportFilePath)
}

Function CleanUpGcpDisk([string]$gcpServiceAccountKeyFile, [string]$cloudDiskName) {
    $gcpServiceAccountKey = Get-Content -Raw -Path $gcpServiceAccountKeyFile | ConvertFrom-Json
    try {
        $disk = Get-GceImage -Name $cloudDiskName -Project $gcpServiceAccountKey.project_id
        if ($disk) {
            Log "Deleting existing disk image $cloudDiskName from GCP $disk"
            Remove-GceImage -Name $cloudDiskName -Project $gcpServiceAccountKey.project_id
        }
    }
    catch {
        Log "No existing disk image $cloudDiskName"
    }
}

Function UploadFromSmbToAws([psobject]$smbConfig, [string]$awsProfileName, [string]$awsRegion, [int]$threads) {
    Log "Uploading disk to AWS"
    $cloudupload = {
        param($SharePath, $Arguments)
        $fullPath = Join-Path -Path $SharePath -ChildPath $Arguments[0]
        Log "Copying '$fullPath' to AWS snapshot"
        $copyArgs = @{File = $fullPath; Threads = $threads}
        if (-not [String]::IsNullOrWhiteSpace($awsProfileName)) { $copyArgs["ProfileName"] = $awsProfileName }
        if (-not [String]::IsNullOrWhiteSpace($awsRegion)) { $copyArgs["Region"] = $awsRegion }
        Copy-ToAwsDisk @copyArgs
    }
    $snapshotId = ExecuteOnSmbShare $cloudupload $smbConfig.SmbCred $smbConfig.ShareUnc @($smbConfig.ExportFilePath)
    Log "Copied disk to AWS snapshot $snapshotId"
    return $snapshotId
}

Function UploadFromSmbToGcp([string]$cloudDiskName, [string]$gcpServiceAccountKeyFile, [psobject]$smbConfig){
    $bucketName = $cloudDiskName.Split('.')[0]
    if (-not (($bucketName.Length -le 63) -and ($bucketName -cmatch '^[a-z]([-a-z0-9]*[a-z0-9])?$'))) {
        ThrowException "Invalid CloudDiskName '$cloudDiskName'. The stem must meet the requirements for Google Cloud image names. See https://cloud.google.com/compute/docs/reference/rest/v1/images."
    }
    Log "Uploading disk to GCP"
    $cloudupload = {
        param($SharePath, $Arguments)
        $fullPath = Join-Path -Path $SharePath -ChildPath $Arguments[0]
        Log "Copying '$fullPath' to '$cloudDiskName' with $gcpServiceAccountKeyFile"
        Copy-ToGcpDisk -File $fullPath -BucketName $bucketName -ServiceAccountKeyFile $gcpServiceAccountKeyFile
    }
    ExecuteOnSmbShare $cloudupload $smbConfig.SmbCred $smbConfig.ShareUnc @($smbConfig.ExportFilePath)
    Log "Copied disk to '$cloudDiskName' in bucket '$bucketName'"
}

Function ExecuteOnSmbShare([ScriptBlock]$scriptblock, [PSCredential]$smbCred, [string]$share, [string[]]$arguments)
{
    # Use this function when the code in the scriptblock is all Powershell.
    $name = "CtxMapping"
    Log "ExecuteOnSmbShareWithCreds as $smbCred on share $share with args $arguments"
    $err = ""
    $drive = New-PSDrive -Name $name -PSProvider "FileSystem" -Root $share -Credential $smbCred -Scope Global -ErrorVariable "err"
    if (!$drive) {
        ThrowException "Failed to map share $share error: $err"
    }
    try {
        Log "Operating on $($drive.Name) with $arguments"
        $output = & $scriptblock -SharePath "$($name):" -Arguments $arguments
    }
    finally {
        $null = Remove-PSDrive -Name $name
    }
    return $output
}

Function InitUploadLog([bool]$overwrite) {
    $Global:UploadLogFile = "$(Get-Location)\Upload.log"
    $timestamp = (Get-Date).ToUniversalTime().ToString("yyyy-MM-ddTHH:mm:ss.ffffZ")
    if ($overwrite) {
        "$($timestamp): New log" | Out-File -FilePath $Global:UploadLogFile
    } else {
        "$($timestamp): New log" | Out-File -FilePath $Global:UploadLogFile -Append
    }
    Write-Host "Logging to $($Global:UploadLogFile)"
}

Function Log([string]$message, [bool]$echoToScreen = $True) {
    if ($echoToScreen) {
        Write-Host $message
    }
    $timestamp = (Get-Date).ToUniversalTime().ToString("yyyy-MM-ddTHH:mm:ss.ffffZ")
    "$($timestamp): $message" | Out-File -FilePath $Global:UploadLogFile -Append
}

Function ThrowException([string]$message, [bool]$echoToScreen = $True) {
    Log $message $echoToScreen
    throw $message
}