PureStorage.AzureNative.Util.ps1

$PluginPrivileges = @(
    'Datastore.AllocateSpace',
    'Datastore.Browse',
    'Datastore.Config',
    'Datastore.Delete',
    'Datastore.DeleteFile',
    'Datastore.FileManagement',
    'Datastore.Move',
    'Datastore.Rename',
    'Datastore.UpdateVirtualMachineFiles',
    'Datastore.UpdateVirtualMachineMetadata',
    'Extension.Register',
    'Extension.Unregister',
    'Extension.Update',
    'Folder.Create',
    'Folder.Delete',
    'Folder.Move',
    'Folder.Rename',
    'Global.CancelTask',
    'Global.ManageCustomFields',
    'Global.SetCustomField',
    'Host.Config.Storage',
    'ScheduledTask.Create',
    'ScheduledTask.Delete',
    'ScheduledTask.Edit',
    'ScheduledTask.Run',
    'Sessions.ValidateSession',
    'StorageProfile.Update',
    'StorageProfile.View',
    'StorageViews.View',
    'StorageViews.ConfigureService',
    'Task.Create',
    'Task.Update',
    'VirtualMachine.Config.AddExistingDisk',
    'VirtualMachine.Config.AddNewDisk',
    'VirtualMachine.Config.AddRemoveDevice',
    'VirtualMachine.Config.RemoveDisk',
    'VirtualMachine.Interact.PowerOff',
    'VirtualMachine.Interact.PowerOn',
    'VirtualMachine.Inventory.Create',
    'VirtualMachine.Inventory.CreateFromExisting',
    'VirtualMachine.Inventory.Delete',
    'VirtualMachine.Inventory.Move',
    'VirtualMachine.Inventory.Register',
    'VirtualMachine.Inventory.Unregister',
    'VirtualMachine.Provisioning.Clone',
    'VirtualMachine.Provisioning.CloneTemplate',
    'VirtualMachine.Provisioning.CreateTemplateFromVM',
    'VirtualMachine.Provisioning.GetVmFiles',
    'VirtualMachine.State.CreateSnapshot',
    'VirtualMachine.State.RemoveSnapshot',
    'VirtualMachine.State.RenameSnapshot',
    'VirtualMachine.State.RevertToSnapshot'
)

# Service account name prefix
$AccountNamePrefix = "psserviceaccount"

# Service account role name
$RoleName = "PureStorageService"

# Timeout for successful Rest call
$WaitTimeSeconds = 3600

$DefaultManagementHostUri = "https://management.azure.com"

# The AVS SKUs that support iSCSI multipath feature. Names are taken from the AVS API /subscriptions/<subscription-id>/providers/Microsoft.AVS/skus?api-version=2024-09-01
$AvsSkusWithIscsiMultipathFeature = @("av20","av36","av36p","av36pt","av36t","av48","av48t","av52","av52t")

# .SYNOPSIS
# Write a message to console with added timestamp.
function PrintLog {
    param (
        [Parameter(Mandatory = $true)]
        [string]$Message,

        [Parameter()]
        [ValidateNotNullOrEmpty()]
        [ValidateSet('INFO', 'DEBUG', 'WARNING', 'ERROR', 'VERBOSE')]
        [string]$Severity = 'INFO'
    )
    $dateTime = (Get-Date -f "yyyy-MM-dd HH:mm:ss")
    $msg = "$dateTime $Severity $Message"
    switch ($Severity) {
        WARNING {
            Write-Warning $msg
        }
        DEBUG {
            Write-Debug $msg
        }
        ERROR {
            Write-Error $msg
        }
        VERBOSE {
            Write-Verbose $msg
        }
        Default {
            Write-Host $msg
        }
    }
}

function New-RandomPassword {
    param (
        [int]$length = 16,
        [int]$specialCharCount = 4
    )
    
    if ($length -lt 8) {
        throw "Password length should be at least 8 characters."
    }

    if ($specialCharCount -gt $length) {
        throw "Number of special characters cannot exceed total length of the password."
    }

    $upperCaseCount = [math]::Ceiling(($length - $specialCharCount) / 3)
    $lowerCaseCount = [math]::Ceiling(($length - $specialCharCount) / 3)
    $numberCount    = ($length - $specialCharCount) - $upperCaseCount - $lowerCaseCount

    $upperCase = 1..$upperCaseCount | ForEach-Object { [char[]]([char]'A'..[char]'Z') | Get-SecureRandom }
    $lowerCase = 1..$lowerCaseCount | ForEach-Object { [char[]]([char]'a'..[char]'z') | Get-SecureRandom }
    $numbers   = 1..$numberCount | ForEach-Object { [char[]]([char]'0'..[char]'9') | Get-SecureRandom }
    $special   = 1..$specialCharCount | ForEach-Object { ([char[]]"!@#$%^&*()_+-=[]{}|;:',.<>?") | Get-SecureRandom }

    $passwordChars = $upperCase + $lowerCase + $numbers + $special
    $password = ($passwordChars | Sort-Object { Get-SecureRandom }) -join ""

    return $password
}

function Get-EncryptedSignature {
    param (
        [Parameter(Mandatory = $true)]
        [string]$Text,
        [Parameter(Mandatory = $true)]
        [string]$PrivateKey
    )

    $hashAlgorithm = [System.Security.Cryptography.HashAlgorithmName]::SHA256

    # Convert the private key from PEM format to an RSA object
    $rsaKey = [System.Security.Cryptography.RSA]::Create()
    $rsaKey.ImportFromPem($PrivateKey)

    $bytesToEncrypt = [System.Text.Encoding]::UTF8.GetBytes($Text)

    $signature = $rsaKey.SignData($bytesToEncrypt, $hashAlgorithm, [System.Security.Cryptography.RSASignaturePadding]::Pkcs1)

    # Convert the encrypted byte array to a base64 string
    $Signature64 = [Convert]::ToBase64String($signature)

    return $Signature64
}

function Test-TextSignarure {
    param (
        [Parameter(Mandatory = $true)]
        [string]$Text,
        [Parameter(Mandatory = $true)]
        [string]$Signature,
        [Parameter(Mandatory = $true)]
        [string]$PublicKey
    )
    
    $key = $PublicKey -replace "-----BEGIN PUBLIC KEY-----", "" -replace "-----END PUBLIC KEY-----", ""
    $key = [Convert]::FromBase64String($key)

    $hashAlgorithm = [System.Security.Cryptography.HashAlgorithmName]::SHA256

    # Convert the private key from PEM format to an RSA object
    $rsaKey = [System.Security.Cryptography.RSA]::Create()
    $rsaKey.ImportSubjectPublicKeyInfo($key, [ref]0)
    
    # Convert the base64-encoded encrypted text to a byte array
    $SignatureBytes = [Convert]::FromBase64String($Signature)

    $TextBytes = [System.Text.Encoding]::UTF8.GetBytes($Text)
    $IsValid =  $rsaKey.VerifyData($TextBytes, $SignatureBytes, $hashAlgorithm, [System.Security.Cryptography.RSASignaturePadding]::Pkcs1)
    
    return $IsValid
}

function Test-RequestDatetimeInUTC {
    param (
        [Parameter(Mandatory = $true)]
        [string]$RequestDatetime,
        [Parameter(Mandatory = $true)]
        [int]$TimeWindowInMinutes
    )
    # Check time window to validate the request
    $requestDatetimeUTC = Get-Date -Date $RequestDatetime
    $currentDatetimeUTC = Get-Date -AsUTC
    if ($requestDatetimeUTC.AddMinutes($TimeWindowInMinutes) -lt $currentDatetimeUTC) {
        throw "Request is outside the time window: $TimeWindowInMinutes minutes"
    }
}

function ConvertTo-EncryptedText {
    param (
        [Parameter(Mandatory = $true)]
        [string]$Text,
        [Parameter(Mandatory = $true)]
        [string]$PublicKey
    )

    # Convert the private key from PEM format to an RSA object
    $rsaKey = [System.Security.Cryptography.RSA]::Create()
    $rsaKey.ImportFromPem($PublicKey)
 
    # Convert the text to a byte array
    $bytesToEncrypt = [System.Text.Encoding]::UTF8.GetBytes($Text)

    # Encrypt the byte array using the key
    $encryptedBytes = $rsaKey.Encrypt($bytesToEncrypt, [System.Security.Cryptography.RSAEncryptionPadding]::Pkcs1)

    # Convert the encrypted byte array to a base64 string
    $encryptedText = [Convert]::ToBase64String($encryptedBytes)

    return $encryptedText
}

function ConvertFrom-EncryptedText {
    param (
        [Parameter(Mandatory = $true)]
        [string]$EncryptedText,
        [Parameter(Mandatory = $true)]
        [string]$PrivateKey
    )

    # Convert the private key from PEM format to an RSA object
    $rsaKey = [System.Security.Cryptography.RSA]::Create()
    $rsaKey.ImportFromPem($PrivateKey)

    # Convert the base64-encoded encrypted text to a byte array
    $encryptedBytes = [Convert]::FromBase64String($EncryptedText)

    # Decrypt the byte array using the key
    $decryptedBytes = $rsaKey.Decrypt($encryptedBytes, [System.Security.Cryptography.RSAEncryptionPadding]::Pkcs1)

    # Convert the decrypted byte array back to a string
    $decryptedText = [System.Text.Encoding]::UTF8.GetString($decryptedBytes)

    return $decryptedText
}

function ConvertFrom-Base64 {
    param (
        [Parameter(Mandatory = $true)]
        [string]$Base64Text
    )

    $bytes = [Convert]::FromBase64String($Base64Text)
    $text = [System.Text.Encoding]::UTF8.GetString($bytes)

    return $text
}

function ConvertTo-Base64 {
    param (
        [Parameter(Mandatory = $true)]
        [string]$Text
    )

    $bytes = [System.Text.Encoding]::UTF8.GetBytes($Text)
    $base64Text = [Convert]::ToBase64String($bytes)

    return $base64Text
}

function CompareAndUpdate-Privileges {
    param (
        [Parameter(Mandatory = $true)]
        [object]$Role,
        [Parameter(Mandatory = $true)]
        [string[]]$RequiredPrivileges
    )
    # These are the default privileges that are always required and cannot be removed
    $DefaultPrivileges = @(
        'System.Anonymous',
        'System.Read',
        'System.View'
    )
    $RequiredPrivileges += $DefaultPrivileges

    $RoleName = $Role.Name
    $RolePrivileges = $Role.PrivilegeList

    # Identify missing and extra privileges
    # Find missing privileges (those in the required list but not in the role)
    $MissingPrivileges = $RequiredPrivileges | Where-Object { $_ -notin $RolePrivileges }
    # Find extra privileges (those in the role but not in the required list)
    $ExtraPrivileges = $RolePrivileges | Where-Object { $_ -notin $RequiredPrivileges }

    $PrivilegesToAdd = @()
    if ($MissingPrivileges) {
        PrintLog "Missing Privileges: $($MissingPrivileges -join ', ')" INFO
        $PrivilegesToAdd += $MissingPrivileges | ForEach-Object { Get-VIPrivilege -Id $_ }
    } else {
        PrintLog "No missing privileges" INFO
    }

    $PrivilegesToRemove = @()
    if ($ExtraPrivileges) {
        PrintLog "Extra Privileges: $($ExtraPrivileges -join ', ')" INFO
        $PrivilegesToRemove += $ExtraPrivileges | ForEach-Object { Get-VIPrivilege -Id $_ }
    } else {
        PrintLog "No extra privileges" INFO
    }

    # Update the role with the missing and extra privileges
    # Apply updates only if needed
    $IsUpdated = $false
    if ($PrivilegesToAdd.Count -gt 0) {
        PrintLog "Adding missing privileges to role '$RoleName'" INFO
        $null = Set-VIRole -Role $Role -AddPrivilege $PrivilegesToAdd
        $IsUpdated = $true
    }

    if ($PrivilegesToRemove.Count -gt 0) {
        PrintLog "Removing extra privileges from role '$RoleName'" INFO
        $null = Set-VIRole -Role $Role -RemovePrivilege $PrivilegesToRemove
        $IsUpdated = $true
    }

    if (-not $IsUpdated) {
        PrintLog "No changes needed for the role: $RoleName" INFO
    }
}

<#
    .SYNOPSIS
    Creates a new service account and assigns it a role with specific privileges.
#>

function _New-AvsServiceAccount {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ServiceInitializationHandleEnc
    )

    # Convert the ServiceInitializationHandleEnc to a JSON object
    $DecodedInitializationHandle = ConvertFrom-Base64 -Base64Text $ServiceInitializationHandleEnc | ConvertFrom-Json
    $Data = $DecodedInitializationHandle.data
    # Signature is ignored as per requirements

    # Convert the data to a JSON object
    $DecodedData = ConvertFrom-Base64 -Base64Text $Data | ConvertFrom-Json

    # The DecodedData is a JSON object with the following structure:
    # {
    # "sddcResourceId": "string",
    # "serviceAccountUsername": "string",
    # "ephemeralPublicKey": "string"
    # }
    $AccountName = $DecodedData.serviceAccountUsername
    $EphemeralPublicKey = $DecodedData.ephemeralPublicKey

    # Validate the prefix of the account name
    if (-not $AccountName.StartsWith($AccountNamePrefix)) {
        throw "The account name must start with '$AccountNamePrefix'"
    }

    # Generate a random password for the service account
    $AccountPassword = New-RandomPassword

    # If the user already exists, update the password
    $User = Get-SsoPersonUser -Domain 'vsphere.local' | Where-Object { $_.Name -eq $AccountName }
    if ($User) {
        PrintLog "User $AccountName already exists, updating the password" WARNING
        $null = Set-SsoPersonUser -User $User -NewPassword $AccountPassword -ErrorAction Stop
    } else {
        $User = New-SsoPersonUser -UserName $AccountName -Password $AccountPassword -Description "Pure Storage Service Account" -ErrorAction Stop
    }
    # Create Role and assign Role to user
    $Role = Get-VIRole -Name $RoleName -ErrorAction SilentlyContinue
    if ($Role) {
        PrintLog "Role $RoleName already exists, checking the role privileges against the required privileges" WARNING
        CompareAndUpdate-Privileges -Role $Role -RequiredPrivileges $PluginPrivileges
    }
    else {
        $Privileges = @()
        foreach ($priv in $PluginPrivileges) {
            PrintLog "Adding privilege: $priv" DEBUG
            $Privileges += Get-VIPrivilege -Id $priv
        }

        $Role = New-VIRole -Name $RoleName -Privilege $Privileges
    }

    $Account = Get-VIAccount -Domain $User.Domain | Where-Object { $_.Id -eq $AccountName }
    if (-not $Account) {
        throw "Failed to create account for user $User"
    }

    $RootFolder = Get-Folder -NoRecursion
    if (-not $RootFolder) {
        throw "Failed to retrieve root folder"
    }
    $RootFolder = $RootFolder | Select-Object -Last 1

    PrintLog "Adding permissions for Account $AccountName on $($RootFolder.Name) with role $RoleName" INFO
    $null = New-VIPermission -Entity $RootFolder -Principal $Account -Role $Role -Propagate $true

    PrintLog "Avs Service Account $AccountName created successfully with role $RoleName" INFO

    # Return the initialization data for the service account
    $vSphereIp = $Account.Server.ServiceUri.Host

    # Encrypt the password with the ephemeral public key
    if ([string]::IsNullOrEmpty($EphemeralPublicKey)) {
        PrintLog "Ephemeral public key is missing from initialization handle" ERROR
        throw "Ephemeral public key is required for password encryption"
    }

    try {
        $EncryptedPassword = ConvertTo-EncryptedText -Text $AccountPassword -PublicKey $EphemeralPublicKey
        PrintLog "Password encrypted successfully with ephemeral public key" INFO
    }
    catch {
        PrintLog "Failed to encrypt the password with error: $_" ERROR
        throw "Failed to encrypt the password with error: $_"
    }

    # Add "encrypted:" prefix to the encrypted password for KAVS to identify it
    $EncryptedPasswordWithPrefix = "encrypted:$EncryptedPassword"

    $InitializationData= @{
        "serviceAccountUsername" = $AccountName
        "serviceAccountPassword" = $EncryptedPasswordWithPrefix
        "vSphereIp" = $vSphereIp
    }

    return $InitializationData
}

<#
    .SYNOPSIS
    Removes Pure Storage AVS service account(s) and the role assigned to them.
#>

function _Remove-AvsServiceAccount
{
    param(
        [Parameter(Mandatory = $false)]
        [string]$Suffix,
        [switch]$DryRun
    )

    # If we want to remove all service accounts, we can set the suffix to "*"
    $AccountName = $AccountNamePrefix + $Suffix
    $users = Get-SsoPersonUser -Domain 'vsphere.local' | Where-Object { $_.Name -like $AccountName }
    if (-not $users) {
        PrintLog "User '$AccountName' not found." WARNING
        return
    }
    PrintLog "Found $($users.Count) user(s) in total" INFO

    foreach ($user in $users) {
        $userName = $user.Name
        PrintLog "Found user: $userName" INFO
        $name = "VSPHERE.LOCAL\" + $user.Name
        $accountPermissions = Get-VIPermission -Principal $name

        # Attempt to remove any permissions the user might have
        if ($accountPermissions) {
            if ($DryRun) {
                PrintLog "Dry run: Removing permissions '$($accountPermissions.Role)' from user $userName" INFO
            } else {
                PrintLog "Removing permissions '$($accountPermissions.Role)' from user $userName" INFO
                try {
                    Remove-VIPermission -Permission $accountPermissions -Confirm:$false
                    PrintLog "Successfully removed permissions for user $userName" INFO
                } catch {
                    PrintLog "Failed to remove permissions for user $userName with error: $_" WARNING
                }
            }
        }

        # Remove the user regardless of role
        if ($DryRun) {
            PrintLog "Dry run: Removing user $userName" INFO
        } else {
            PrintLog "Removing user $userName" INFO
            try {
                Remove-SsoPersonUser -User $user
                PrintLog "Successfully removed user $userName." INFO
            } catch {
                throw "Failed to remove user $userName with error: $_"
            }
        }

        # Check and clean up role if it's unused
        $role = Get-VIRole -Name $RoleName
        if ($role) {
            $remainingRoleAssignments = Get-VIPermission | Where-Object { $_.Role -eq $RoleName }
            if ($remainingRoleAssignments.Count -eq 0) {
                if ($DryRun) {
                    PrintLog "Dry run: Removing unused role $RoleName" INFO
                } else {
                    PrintLog "Removing unused role $RoleName" INFO
                    try {
                        Remove-VIRole -Role $role -Force -Confirm:$false
                        PrintLog "Successfully removed role $RoleName" INFO
                    } catch {
                        throw "Failed to remove role $RoleName with error: $_"
                    }
                }
            } else {
                PrintLog "Role $RoleName still has other accounts assigned" WARNING
            }
        } else {
            PrintLog "Role $RoleName not found" WARNING
        }
    }
}

function ConvertTo-HttpHeaders {
    param (
        [Parameter(Mandatory = $true)]
        [System.Collections.Generic.Dictionary[string, System.Collections.Generic.IEnumerable[string]]]$Headers
    )
    # Create a new HttpRequestMessage to construct HttpHeaders
    $httpRequestMessage = [System.Net.Http.HttpRequestMessage]::new()

    foreach ($key in $Headers.Keys) {
        foreach ($value in $Headers[$key]) {
            $null = $httpRequestMessage.Headers.TryAddWithoutValidation($key, $value)
        }
    }

    $result = $httpRequestMessage.Headers
    Write-Output -NoEnumerate $result
}

<#
    .SYNOPSIS
    Invoke an arm request to Azure REST API and convert the response to an AzRest response object.
#>

function Invoke-ARM{
    param(
        [Parameter(Mandatory = $true)]
        [string]$Uri,
        [Parameter(Mandatory = $true)]
        [string]$Method,
        [Parameter(Mandatory = $false)]
        [string]$Body,
        [Parameter(Mandatory = $false)]
        [hashtable]$RequestHeaders = @{}
    )
    $content = Invoke-ARMRequest -Uri $Uri -Method $Method -Body $Body -Headers $RequestHeaders `
                -StatusCodeVariable "statusCode" -ResponseHeadersVariable "headers"

    if (-not $statusCode) {
        PrintLog "Status code is null. Failed to retrieve status code from ARM request" WARNING
    }
    if (-not $headers) {
        PrintLog "Response headers is null. Failed to retrieve response headers from ARM request" WARNING
    }

    $httpHeaders = ConvertTo-HttpHeaders -Headers $headers

    if (-not $httpHeaders) {
        PrintLog "Failed to convert response headers to HttpHeaders" WARNING
    }

    # Structure the output to match Invoke-AzRest response
    $responseObject = [pscustomobject]@{
        Content = $content
        Headers = $httpHeaders
        StatusCode = $statusCode
    }

    Write-Output -NoEnumerate $responseObject
}

function Invoke-AzureRestCallWithVerificationAndRetries {
    param (
        [Parameter(Mandatory = $true)]
        [ValidateSet("GET", "POST", "PUT", "DELETE")]
        [string]$Method,
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string]$Uri,
        [Parameter(Mandatory = $false)]
        [hashtable]$PayLoad = @{},
        [Parameter(Mandatory = $false)]
        [hashtable]$RequestHeaders = @{},
        [Parameter(Mandatory = $false)]
        [int]$MaxRetries = 3,
        [int]$RetryInterval = 1
    )
    if ($Payload.Count -eq 0) {
        # When payload is empty, set it to an empty string to avoid the Azure GET request failure.
        $jsonPayload = ""
    } else {
        $jsonPayLoad = $PayLoad | ConvertTo-Json -Depth 10
    }

    PrintLog "Request URI: $Uri" DEBUG
    PrintLog "Request BODY: $jsonPayLoad" DEBUG

    $attempt = 0
    while ($attempt -lt $MaxRetries) {
        try {
            $attempt++
            $res = Invoke-ARM -Uri $Uri -Method $Method -Body $jsonPayLoad -RequestHeaders $RequestHeaders

            if ($res.StatusCode -ge 200 -and $res.StatusCode -lt 300) {
                return $res
            }

            $errorMsg = "Request failed with status code: $( $res.StatusCode ) and content: $( $res.Content )"
            PrintLog $errorMsg ERROR
            throw $errorMsg
        }
        catch {
            $errorMsg = $_.Exception.Message
            if ($attempt -lt $MaxRetries)
            {
                PrintLog "Attempt $attempt failed. Retrying in $RetryInterval seconds..." DEBUG
                Start-Sleep -Seconds $RetryInterval
            } else {
                $msg = "Maximum retry attempts reached: $errorMsg"
                PrintLog $msg ERROR
                throw $msg
            }
        }
    }
}

function WaitForSuccessStatus {
    param(
        [Parameter(Mandatory = $true)]
        [System.Net.Http.Headers.HttpHeaders]$Headers,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $false)]
        [int]$WaitTimeoutSeconds,
        [Parameter(Mandatory = $false)]
        [int]$CheckIntervalSeconds = 30
    )
    $asyncOpUrl = $Headers.GetValues("Azure-AsyncOperation")
    $correlationId = $Headers.GetValues("x-ms-correlation-request-id")
    $requestHeaders = Get-RequestHeader -CorrelationId $CorrelationId
    PrintLog "'x-ms-correlation-request-id' of the request: '$correlationId'" DEBUG
    PrintLog "Op URL: $asyncOpUrl" DEBUG

    PrintLog "Querying the operation status for up to $WaitTimeoutSeconds seconds to complete"
    $endTime = (Get-Date).AddSeconds($WaitTimeoutSeconds)
    while ((Get-Date) -lt $endTime) {
        $content = (Invoke-AzureRestCallWithVerificationAndRetries -Method 'GET' -Uri "$asyncOpUrl" -RequestHeaders $requestHeaders).Content
        $opStatus = $content.status
        PrintLog "Operation status: '$opStatus'" DEBUG
        if ($opStatus -eq "Succeeded") {
            return ""
        }
        elseif ($opStatus -ne "Accepted" -And $opStatus -ne "Deleting" -And $opStatus -ne "Pending" -And $opStatus -ne "Running") {
            $msg = "OP status for '$correlationId' correlation ID: '$opStatus'"
            PrintLog $msg ERROR
            if ($opStatus -eq "Failed") {
                $errorMessage = "{0}" -f $content.error.message
                PrintLog "Error message: $errorMessage" ERROR
                return $errorMessage
            }
            $msg = "Unexpected OP status"
            PrintLog $msg ERROR
            return $msg
        }
        Start-Sleep -Seconds $CheckIntervalSeconds
    }
    $msg = "Timed out waiting for successful operation status, correlation ID: '$correlationId'"
    PrintLog $msg ERROR
    return $msg
}

function Get-RequestHeader {
    param (
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId
    )
    $requestHeaders = @{
        'Content-Type' = 'application/json; charset=utf-8'
    }

    if ($CorrelationId -and $CorrelationId -ne '') {
        $requestHeaders['x-krypton-correlation-id'] = $CorrelationId
    }

    return $requestHeaders
}

function Invoke-EnableAvsConnection {
    param (
        [Parameter(Mandatory = $true)]
        [string]$StoragePoolResourceId,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $true)]
        [string]$PureStorageBlockApiVersion
    )
    if ($null -ne $SddcResourceId) {
        PrintLog "SDDC resource ID: $SddcResourceId" DEBUG
        $payload = @{
            'sddcResourceId' = $SddcResourceId
        }
    } else {
        throw "SDDC resource ID is empty"
    }

    $requestHeaders = Get-RequestHeader -CorrelationId $CorrelationId
    $uri = $DefaultManagementHostUri + $StoragePoolResourceId + "/enableAvsConnection?api-version=$PureStorageBlockApiVersion"

    $res = Invoke-AzureRestCallWithVerificationAndRetries -Method 'POST' -Uri $uri -PayLoad $payload -RequestHeaders $requestHeaders
    if ($null -eq $res) {
        $msg =  "Azure rest call failed. No response received."
        PrintLog $msg ERROR
        return $msg
    }

    return WaitForSuccessStatus -Headers $res.Headers -WaitTimeoutSeconds $WaitTimeSeconds -CorrelationId $CorrelationId
}

function Get-AvsConnection {
    param (
        [Parameter(Mandatory = $true)]
        [string]$StoragePoolResourceId,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $true)]
        [string]$PureStorageBlockApiVersion
    )
    $requestHeaders = Get-RequestHeader -CorrelationId $CorrelationId
    $uri = $DefaultManagementHostUri + $StoragePoolResourceId + "/getAvsConnection?api-version=$PureStorageBlockApiVersion"

    return Invoke-AzureRestCallWithVerificationAndRetries -Method 'POST' -Uri $uri -RequestHeaders $requestHeaders `
        | Select-Object -ExpandProperty Content
}

function Invoke-FinalizeAvsConnection {
    param (
        [Parameter(Mandatory = $true)]
        [hashtable]$ServiceInitializationData,
        [Parameter(Mandatory = $true)]
        [string]$StoragePoolResourceId,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $true)]
        [string]$PureStorageBlockApiVersion
    )

    $payload = @{
        'serviceInitializationDataEnc' = ""
        'serviceInitializationData' = $ServiceInitializationData
    }

    $requestHeaders = Get-RequestHeader -CorrelationId $CorrelationId
    $uri = $DefaultManagementHostUri + $StoragePoolResourceId + "/finalizeAvsConnection?api-version=$PureStorageBlockApiVersion"

    $res = Invoke-AzureRestCallWithVerificationAndRetries -Method 'POST' -Uri $uri -PayLoad $payload -RequestHeaders $requestHeaders
    if ($null -eq $res) {
        $msg =  "Azure rest call failed. No response received."
        PrintLog $msg ERROR
        return $msg
    }

    return WaitForSuccessStatus -Headers $res.Headers -WaitTimeoutSeconds $WaitTimeSeconds -CorrelationId $CorrelationId
}

function Invoke-DisableAvsConnection {
    param (
        [Parameter(Mandatory = $true)]
        [string]$StoragePoolResourceId,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $true)]
        [string]$PureStorageBlockApiVersion
    )

    $requestHeaders = Get-RequestHeader -CorrelationId $CorrelationId
    $uri = $DefaultManagementHostUri + $StoragePoolResourceId + "/disableAvsConnection?api-version=$PureStorageBlockApiVersion"

    $res = Invoke-AzureRestCallWithVerificationAndRetries -Method 'POST' -Uri $uri -RequestHeaders $requestHeaders
    if ($null -eq $res) {
        $msg =  "Azure rest call failed. No response received."
        PrintLog $msg ERROR
        return $msg
    }

    return WaitForSuccessStatus -Headers $res.Headers -WaitTimeoutSeconds $WaitTimeSeconds -CorrelationId $CorrelationId
}

function Get-AVSInfo {
    param (
        [Parameter(Mandatory = $false)]
        [string]$AvsApiVersion = "2024-09-01"
    )

    $uri = $DefaultManagementHostUri + $SddcResourceId + "?api-version=$AvsApiVersion"
    $res = Invoke-AzureRestCallWithVerificationAndRetries -Method 'GET' -Uri  $uri 

    return $res.Content
}

function Get-AvsISCSIPaths {
    param (
        [Parameter(Mandatory = $false)]
        [string]$AvsApiVersion = "2024-09-01"
    )

    $uri = $DefaultManagementHostUri + $SddcResourceId + "/iscsiPaths?api-version=$AvsApiVersion"
    $res = Invoke-AzureRestCallWithVerificationAndRetries -Uri $uri -Method 'GET' 

    return $res.Content.value
}

function Invoke-IsAvsMultipathEnabled {
    param (
        [Parameter(Mandatory = $false)]
        [string]$AvsApiVersion = "2024-09-01"
    )
    PrintLog "Start AVS SKU validation for iSCSI multipath feature." INFO
    $AvsInfo = Get-AVSInfo -AvsApiVersion $AvsApiVersion
    if ($AvsSkusWithIscsiMultipathFeature -notcontains $AvsInfo.sku.name) {
        PrintLog "AVS SKU '$($AvsInfo.sku.name)' does not support iSCSI multipath feature. Skipping validation." INFO
        return
    }

    PrintLog "Start AVS SKU validation." INFO

    try {
        $iscsiPaths = Get-AvsISCSIPaths -AvsApiVersion $AvsApiVersion
    } catch {
        throw "Failed to validate iSCSI Multipath: $_"
    }

    # We expect only one path to be presented
    if ($iscsiPaths.Count -ne 1) {
        throw "Failed to validate iSCSI Multipath: Invalid amount of iSCSI paths found for AVS SKU. Expected 1, found $($iscsiPaths.Count)"
    }

    $defaultPath = $iscsiPaths[0]
    if ($defaultPath.properties.provisioningState -notlike "Succeeded") {
        throw "Failed to validate iSCSI Multipath: iSCSI path $($defaultPath.name) is not in Succeeded state. Current state: $($defaultPath.properties.provisioningState)"
    }
}

function Invoke-HandleRuncommandError {
    param (
        [Parameter(Mandatory = $true)]
        [string]$ErrorMessage,
        [Parameter(Mandatory = $true)]
        [string]$StoragePoolResourceId,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $true)]
        [string]$PureStorageBlockApiVersion
    )
    PrintLog "STEP: Handling RunCommand Error" INFO
    PrintLog "ErrorMessage: $ErrorMessage" INFO

    $errorReportingCodeSA = "error"

    $encodedErrorMessageSA = ConvertTo-Base64 -Text $ErrorMessage
    $initializationData = @{
        "serviceAccountUsername" = ""
        "serviceAccountPassword" = ""
        "vSphereIp" = $errorReportingCodeSA
        "vSphereCertificate" = $encodedErrorMessageSA
    }

    $err = Invoke-FinalizeAvsConnection -ServiceInitializationData $initializationData `
                -StoragePoolResourceId $StoragePoolResourceId `
                -CorrelationId $CorrelationId `
                -PureStorageBlockApiVersion $PureStorageBlockApiVersion
    if ($err) {
        throw $err
    }
}

<#
    .SYNOPSIS
    Check if the datastore is a Pure Storage VVOL datastore.
#>

function Is-PSVvolDatastore {
    param (
        [Parameter(Mandatory = $true)]
        [string]$DatastoreName
    )
    $datastore = Get-Datastore -Name $DatastoreName
    if ($null -eq $datastore) {
        PrintLog "Datastore $DatastoreName not found" WARNING
        return $false
    }

    # Check if the datastore is a VVOL datastore as we only support VVOL in AVS
    if ($datastore.Type -eq "VVOL") {
        # Check the Storage Provider of the datastore
        $dsSPInfo = $datastore.ExtensionData.Info.VvolDS.VasaProviderInfo
        $arrayId = $dsSPInfo.ArrayState.ArrayId
        PrintLog "arrayId: $arrayId" DEBUG
        If ($arrayId -like "com.purestorage*") {
            PrintLog "Datastore $DatastoreName is a Pure Storage VVOL datastore" INFO
            return $true
        }
    }
    return $false
}

function Remove-VvolDatastore {
    param (
        [switch]$DryRun
    )
    # Retrieve the list of clusters
    $clusters = Get-Cluster
    if (-not $clusters) {
        PrintLog "No clusters found" WARNING
        return
    }

    # Iterate through each cluster
    foreach ($cluster in $clusters) {
        # Get the list of datastores in the cluster
        PrintLog "Retrieving datastores for cluster $($cluster.Name)..." INFO
        $datastores = Get-Datastore -RelatedObject $cluster
        if (-not $datastores) {
            PrintLog "No datastores found in cluster $($cluster.Name)" INFO
            continue
        }

        # Iterate through each datastore
        foreach ($datastore in $datastores) {
            # Check if the datastore is a PS VVOL datastore
            $isPSVvol = Is-PSVvolDatastore -DatastoreName $datastore.Name
            if ($isPSVvol) {
                # Check if there are any VMs using the datastore
                $vms = $datastore | Get-VM
                if ($vms) {
                    throw "Cannot remove the datastore $($datastore.Name) as it is in use by VMs. Please remove the VMs first"
                }
                PrintLog "Removing VVOL datastore: $($datastore.Name)..." INFO
                $vmHosts = $cluster | Get-VMHost
                # We need to loop through Esxi to unmount the datastore from all of hosts
                foreach ($esxi in $vmHosts) {
                    # Unmount the datastore from the host
                    if ($DryRun) {
                        PrintLog "Dry run: Unmounting datastore $($datastore.Name) from host $($esxi.Name)..." INFO
                    } else {
                        $datastoreSystem = Get-View -Id $esxi.ExtensionData.ConfigManager.DatastoreSystem
                        PrintLog "Unmounting datastore $($datastore.Name) from host $($esxi.Name)..." INFO
                        $datastoreSystem.RemoveDatastore($datastore.ExtensionData.MoRef)
                    }
                }
            } else {
                PrintLog "Datastore $($datastore.Name) is not a Pure Storage VVOL datastore, skipping..." INFO
            }
        }
    }
}

function Remove-PSRemotePlugin {
    param (
        [switch]$DryRun
    )
    # Get the list of registered PS plugins
    $services = Get-View 'ServiceInstance'
    $extensionMgr = Get-View $services.Content.ExtensionManager
    $psExtensions = $extensionMgr.ExtensionList | Where-Object { $_.Key -like "com.purestorage.integrations.vmware.pureplugin*" }

    foreach ($psExtension in $psExtensions) {
        $pluginKey = $psExtension.Key
        $pluginVersion = $psExtension.Version
        if ($DryRun) {
            PrintLog "Dry run: Unregistering Pure Storage plugin $pluginKey with version $pluginVersion..." INFO
        } else {
            PrintLog "Unregistering Pure Storage plugin $pluginKey with version $pluginVersion..." INFO
            $extensionMgr.UnregisterExtension($pluginKey)
        }
    }
}

function Remove-PSStorageProvider {
    param (
        [switch]$DryRun
    )
    $PSStorageProviders = Get-VasaProvider | Where-Object {$_.Namespace -eq "com.purestorage"}
    if ($null -eq $PSStorageProviders) {
        PrintLog "No Pure Storage Storage Provider registered" INFO
    } else {
        foreach ($PSStorageProvider in $PSStorageProviders) {
            if ($DryRun) {
                PrintLog "Dry run: Unregistering Storage Provider $($PSStorageProvider.Name)..." INFO
            } else {
                PrintLog "Unregistering Storage Provider $($PSStorageProvider.Name)..." INFO
                Remove-VasaProvider -Provider $PSStorageProvider -Confirm:$false
            }
        }
    }
}

function Remove-IscsiStaticTargets {
    param (
        [switch]$DryRun
    )
    # Retrieve the list of clusters
    $clusters = Get-Cluster
    if (-not $clusters) {
        PrintLog "No clusters found" WARNING
        return
    }

    foreach ($cluster in $clusters) {
        $targetChanged = $false
        $vmHosts = $cluster | Get-VMHost
        if (-not $vmHosts) {
            PrintLog "No hosts found in cluster $($cluster.Name)" INFO
            continue
        }

        # Remove iSCSI ip address from static discovery from all of hosts if there's a match
        # HBAs (Host Bus Adapters)
        $hbas = $vmHosts | Get-VMHostHba -Type iScsi
        foreach ($hba in $hbas) {
            $targets = $hba | Get-IScsiHbaTarget | Where-Object {($_.Type -eq "Static")}
            foreach ($target in $targets) {
                # Check the iSCSI name of the target to see if it's a Pure target
                if ($target.IscsiName -like "*com.purestorage*") {
                    If ($DryRun) {
                        PrintLog "Dry run: Removing iSCSI target $target from host $($hba.VMHost.Name)..." INFO
                    } else {
                        PrintLog "Removing iSCSI target $target from host $($hba.VMHost.Name)..." INFO
                        $target | Remove-IScsiHbaTarget -Confirm:$false
                        $targetChanged = $true
                    }
                }
            }
        }

        if ($targetChanged) {
            # Rescan after removing the iSCSI targets
            PrintLog "Rescanning storage..." INFO
            $cluster | Get-VMHost | Get-VMHostStorage -RescanAllHba -RescanVMFS | Out-Null
        }
    }
}

function Remove-PSStoragePolicy {
    param (
        [switch]$DryRun
    )
    # Get all storage policies
    $policies = Get-SpbmStoragePolicy

    if ($null -eq $policies) {
        PrintLog "No storage policies found" WARNING
        return
    }

    # Filter the policies to find Pure Storage policies
    foreach ($policy in $policies) {
        # Check the rule sets to determine if it is a Pure Storage policy
        $ruleSets = $policy.AnyOfRuleSets
        $isPSPolicy = $false
        foreach ($rule in $ruleSets) {
            $capabilityName = $rule.AllOfRules.Capability.Name
            # E.g. com.purestorage.storage.policy.LocalSnapshotProtection.LocalSnapshotInterval
            if ($capabilityName -like "com.purestorage.storage.policy*") {
                PrintLog "Found Pure Storage storage policy: $($policy.Name)" INFO
                $isPSPolicy = $true
                break
            }
        }
        if ($isPSPolicy) {
            if ($DryRun) {
                PrintLog "Dry run: Removing Pure Storage storage policy $($policy.Name)..." INFO
            } else {
                PrintLog "Removing Pure Storage storage policy $($policy.Name)..." INFO
                Remove-SpbmStoragePolicy -StoragePolicy $policy -Confirm:$false
            }
        }
    }
}