PureStorage.AzureNative.Tools.psm1

. $PSScriptRoot/PureStorage.AzureNative.Util.ps1

# AVS resource ID for the dev environment
$AvsResourceIdDev = "/subscriptions/96798833-9949-4913-a555-b0f2de70a444/resourceGroups/rg-krypton-sddc-network-dev-eastus/providers/Microsoft.AVS/privateClouds/krypton-private-cloud-dev"
$AvsResourceIdQA = "/subscriptions/de8f0119-1ff4-4aa1-b5ee-7be650f2d750/resourceGroups/rg-krypton-avs-qa-eastus/providers/Microsoft.AVS/privateClouds/krypton-avs-qa-eastus"

# Time window for the request (set to 60 minutes for now, can be adjusted as discussed)
$TimeWindowInMinutes = 60

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

.DESCRIPTION
The New-AvsServiceAccount function creates a new service account with the specified name and password. It then creates a role named 'PureStorageService'
if it doesn't already exist, and assigns the role to the service account. The function also adds permissions for the service account on all VM hosts.

.PARAMETER InitializationHandle
The InitializationHandle is a base64 encoded JSON object that contains the following fields:
{
     "data": "<Base64 encoded InitializationHandle>",
     "signature": "<Signature>"
}

.EXAMPLE
$InitializationDataEnc = New-AvsServiceAccount -InitializationHandle "eyJkYXRh"

This example decodes the InitializationHandle and validates the signature first.
It then creates a new service account with the name defined in data.serviceAccountUsername wiht random password.
It assigns the 'PureStorageService' role to the service account and adds permissions for the service account on all VM hosts.
It returns an encrypted initialization data which is base64 encoded that contains the service account username, password and vSphere IP.

#>

function New-AvsServiceAccount {
    [CmdletBinding()]
    [AVSAttribute(10, UpdatesSDDC = $false, AutomationOnly = $true)]
    param(
        [Parameter(Mandatory = $true)]
        [string]$InitializationHandle
    )
    # TODO: this function will be deprecated once we move to OBO fully.
    # Convert the InitializationHandle to a JSON object
    $DecodedInitializationHandle = ConvertFrom-Base64 -Base64Text $InitializationHandle | ConvertFrom-Json
    $Data = $DecodedInitializationHandle.data
    $Signature = $DecodedInitializationHandle.signature

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

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

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

    # Make sure user doesn't try to use the data with a different SDDC
    if ($env:SddcResourceId -and $env:SddcResourceId -ne $SddcResourceId) {
        throw "The SDDC resource ID in the request does not match the current SDDC resource ID"
    }

    # Determine which keys to use based on the SDDC resource ID and test signature validity
    $keysToTest = @()
    if ($SddcResourceId -eq $AvsResourceIdDev -or $SddcResourceId -eq $AvsResourceIdQA) {
        $keysToTest += "$PSScriptRoot/avs_public_key_dev.pem"
    }
    if (-not [string]::IsNullOrEmpty($SddcResourceId)) {
        $keysToTest += "$PSScriptRoot/avs_public_key_prod.pem"
    } else {
        # Default to local test environment when AVS is not set
        $keysToTest += "$PSScriptRoot/avs_public_key_test.pem"
    }

    # Test signature validity with each key in the list
    $isValidSignature = $false
    foreach ($keyPath in $keysToTest) {
        $publicKey = Get-Content $keyPath -Raw
        $isValidSignature = Test-TextSignarure -Text $Data -Signature $Signature -PublicKey $publicKey
        if ($isValidSignature) {
            break
        }
    }

    if (-not $IsValidSignature) {
        throw "The data signature is not valid"
    }


    # Make sure the request date is in UTC
    if ($RequestDatetime.Kind -ne [System.DateTimeKind]::Utc) {
        throw "Request datetime must be in UTC"
    }
    # Validate the request datetime is within the time window
    Test-RequestDatetimeInUTC -RequestDatetime $RequestDatetime -TimeWindowInMinutes $TimeWindowInMinutes

    # 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) {
        Write-Warning "User $AccountName already exists, updating the password"
        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) {
        Write-Warning "Role $RoleName already exists, checking the role privileges against the required privileges"
        CompareAndUpdate-Privileges -Role $Role -RequiredPrivileges $PluginPrivileges
    }
    else {
        $Privileges = @()
        foreach ($priv in $PluginPrivileges) {
            Write-Debug "Adding privilege: $priv"
            $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"
    }

    Write-Host "Adding permissions for Account $AccountName on $($RootFolder.Name) with Role $RoleName"
    New-VIPermission -Entity $RootFolder -Principal $Account -Role $Role -Propagate $true

    $vSphereIp = $Account.Server.ServiceUri.Host

    $InitializationData= @{
        "serviceAccountUsername" = $AccountName
        "serviceAccountPassword" = $AccountPassword
        "vSphereIp" = $vSphereIp
    } | ConvertTo-Json

    try {
        $InitializationDataEnc = ConvertTo-EncryptedText -Text $InitializationData -PublicKey $EphemeralPublicKey
    }
    catch {
        throw "Failed to encrypt the initialization data with error: $_"
    }

    $NamedOutputs = @{}
    $NamedOutputs["InitializationDataEnc"] = $InitializationDataEnc
    Set-Variable -Name NamedOutputs -Value $NamedOutputs -Scope Global
}

<#
.SYNOPSIS
Removes a service account and the role assigned to it.

.DESCRIPTION
The Remove-AvsServiceAccount function removes the service account with "psserviceaccount" as the name prefix and an optional suffix in dev enviroment.
It removes the account only when it has the "PureStorageService" role assigned to it.
It also removes the role "PureStorageService" after the removal of the account.

.PARAMETER Suffix
The suffix of the account name which is only needed in dev enviroment and is empty in production.

.EXAMPLE
Remove-AvsServiceAccount -Suffix "1234"

This example tries to remove a service account named "psserviceaccount1234".
It checks if the account has the "PureStorageService" role assigned to it and removes the account and the role if it does.

#>

function Remove-AvsServiceAccount
{
    [CmdletBinding()]
    [AVSAttribute(10, UpdatesSDDC = $false, AutomationOnly = $true)]
    param(
        [Parameter(Mandatory = $false)]
        [string]$Suffix
    )
    # TODO: this function will be deprecated once we move to OBO fully.
    $AccountName = $AccountNamePrefix + $Suffix

    $User = Get-SsoPersonUser -Domain 'vsphere.local' | Where-Object { $_.Name -eq $AccountName }
    if ($User) {
        # Get the roles assigned to the user
        $name = "VSPHERE.LOCAL\" + $User.Name
        $accountPermissions = Get-VIPermission -Principal $name
        # Check if the user has the PureStorageService role
        $hasRole = $false
        foreach ($permission in $accountPermissions) {
            if ($permission.Role -eq $RoleName) {
                $hasRole = $true
                break
            }
        }
        if ($hasRole) {
            Write-Host "Removing user $AccountName"
            try {
                Remove-SsoPersonUser -User $User
                # Removes the permission (user and role association)
                Remove-VIPermission -Permission $accountPermissions -Confirm:$false
                Write-Output "Removed user $name from the role $RoleName."
            }
            catch {
                throw "Failed to remove user $AccountName with error: $_"
            }

            $Role = Get-VIRole -Name $RoleName
            if ($Role) {
                # Filter permissions by role name
                $Permissions = Get-VIPermission
                $accountWithRole = $Permissions | Where-Object { $_.Role -eq $RoleName }
                # Remove the PureStorageService role if there is no account assigned to it
                if ($accountWithRole.Count -eq 0) {
                    Write-Host "Removing role $Role "
                    try {
                        Remove-VIRole -Role $Role -Force -Confirm:$false
                    }
                    catch {
                        throw "Failed to remove role $Role with error: $_"
                    }
                } else {
                    Write-Warning "Role $RoleName still has other accounts assigned to it"
                }
            }
            else {
                Write-Warning "Failed to find role $Role"
            }

        } else {
            Write-Warning "This command is only supposed to remove the account with the role $RoleName"
        }
    }
    else {
        Write-Warning "Failed to find user $AccountName"
    }
}

<#
    .SYNOPSIS
    Establish or re-establish the Storagepool-AVS connection.

    .DESCRIPTION
    This function automates the end-to-end process of integrating an Azure Storage Pool with Azure VMware Solution (AVS).
    It performs a series of steps to securely enable the AVS connection, provision required service accounts, and finalize the integration.

    .PARAMETER StoragePoolResourceId
    The resource ID of the storage pool to call enable, finalize and get avs connection endpoint.

    .PARAMETER CorrelationId
    Optional. The Pure Storage Cloud correlation ID for the AVS connection. It's generated by Azure and used to track the request.

    .PARAMETER PureStorageBlockApiVersion
    Optional. The API version for the PureStorage.Block endpoint. (default: '2024-10-01-preview')

    .EXAMPLE
    Connect-StoragepoolToAvs -StoragePoolResourceId "1234" -CorrelationId "5678" -PureStorageBlockApiVersion "2024-10-01-preview"

    .INPUTS
    Storagepool resource Id, Pure Storage Cloud correlation ID, PureStorage.Block API version

    .OUTPUTS
    None
#>

function Connect-StoragepoolToAvs {
    [CmdletBinding()]
    [AVSAttribute(60, UpdatesSDDC = $false, AutomationOnly = $false)]
    param(
        [Parameter(Mandatory = $true)]
        [string]$StoragePoolResourceId,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $false)]
        [string]$PureStorageBlockApiVersion = '2024-10-01-preview'
    )
    # STEP 1. Enable AVS Connection
    PrintLog "STEP 1: Enabling AVS connection" INFO
    if ($err = Invoke-EnableAvsConnection -StoragePoolResourceId $StoragePoolResourceId -CorrelationId $CorrelationId -PureStorageBlockApiVersion $PureStorageBlockApiVersion) {
        throw $err
    }

    $res = Get-AvsConnection -StoragePoolResourceId $StoragePoolResourceId `
                -CorrelationId $CorrelationId `
                -PureStorageBlockApiVersion $PureStorageBlockApiVersion

    if ($res.serviceInitializationCompleted) {
        PrintLog "AVS connection initialization is already completed" INFO
        return
    }

    $serviceInitializationHandleEnc = $res.serviceInitializationHandleEnc
    if ($null -eq $serviceInitializationHandleEnc -or $serviceInitializationHandleEnc -eq "") {
        throw "Failed to enable AVS connection: serviceInitializationHandleEnc is missing"
    }

    # STEP 2. Call Create Service Account Run Command
    $errorMessageSA = $null
    $errorReportingCodeSA = "error"
    $initializationData = @{}
    PrintLog "STEP 2: Creating a new AVS service account" INFO
    try {
        $initializationData = _New-AvsServiceAccount -ServiceInitializationHandleEnc $serviceInitializationHandleEnc
    } catch {
        $errorMessageSA = "AVS service account creation failed: $($_.Exception.Message)"
        
        Invoke-HandleRuncommandError -ErrorMessage $errorMessageSA `
                -StoragePoolResourceId $StoragePoolResourceId `
                -CorrelationId $CorrelationId `
                -PureStorageBlockApiVersion $PureStorageBlockApiVersion
        throw "Failed to create AVS service account with error: $errorMessageSA"
    }

    # STEP 3. Finalize AVS Connection
    PrintLog "STEP 3: Finalizing Avs Connection" INFO

    $err = Invoke-FinalizeAvsConnection -ServiceInitializationData $initializationData `
                -StoragePoolResourceId $StoragePoolResourceId `
                -CorrelationId $CorrelationId `
                -PureStorageBlockApiVersion $PureStorageBlockApiVersion
    if ($err) {
        $msg = "Failed to finalize AVS connection with error: $err"
        throw $msg
    }

    PrintLog "AVS connection is established successfully" INFO
}

<#
    .SYNOPSIS
    Disable or remove an existing Storagepool–AVS connection in a controlled manner.

    .DESCRIPTION
    This function automates the end-to-end process of removing an Azure Storage Pool from Azure VMware Solution (AVS).
    It securely disables the AVS connection by calling the liftr API.

    .PARAMETER StoragePoolResourceId
    The resource ID of the storage pool to call disable avs connection endpoint.

    .PARAMETER CorrelationId
    Optional. The Pure Storage Cloud correlation ID for the AVS connection. It's generated by Azure and used to track the request.

    .PARAMETER PureStorageBlockApiVersion
    Optional. The API version for the PureStorage.Block endpoint. (default: '2024-10-01-preview')

    .EXAMPLE
    Disconnect-StoragepoolFromAvs -StoragePoolResourceId "1234" -CorrelationId "5678" -PureStorageBlockApiVersion "2024-10-01-preview"

    .INPUTS
    Storagepool resource Id, Pure Storage Cloud correlation ID, PureStorage.Block API version

    .OUTPUTS
    None
#>

Function Disconnect-StoragepoolFromAvs {
    [CmdletBinding()]
    [AVSAttribute(60, UpdatesSDDC = $false, AutomationOnly = $false)]
    param(
        [Parameter(Mandatory = $true)]
        [string]$StoragePoolResourceId,
        [Parameter(Mandatory = $false)]
        [string]$CorrelationId,
        [Parameter(Mandatory = $false)]
        [string]$PureStorageBlockApiVersion = '2024-10-01-preview'
    )

    # Disable AVS Connection
    PrintLog "STEP: Disabling AVS connection" INFO
    $err = Invoke-DisableAvsConnection -StoragePoolResourceId $StoragePoolResourceId -CorrelationId $CorrelationId -PureStorageBlockApiVersion $PureStorageBlockApiVersion
    if ($err) {
        $msg = "Failed to disable AVS connection with error: $err"
        throw $msg
    }

    PrintLog "Storage pool is disconnected from AVS successfully" INFO
}

<#
    .SYNOPSIS
    Remove all Pure Storage Cloud resources from AVS.

    .DESCRIPTION
    This function removes ALL Pure Storage Cloud resources and configurations from Azure VMware Solution (AVS).
    Use this command when the standard Disconnect-StoragepoolFromAvs workflow fails or when the connection is broken.

    WARNING: This is a destructive operation that will remove:
    - All Pure Storage Cloud vVol datastores
    - All Pure Storage Cloud iSCSI static targets
    - The Pure Storage Cloud vSphere Client plugin
    - All Pure Storage Cloud VASA providers
    - All Pure Storage Cloud AVS service accounts
    - All Pure Storage Cloud storage policies

    By default, DryRun is enabled ($true) to preview changes without making them.
    Set DryRun to $false to perform the actual cleanup.

    .PARAMETER DryRun
    When $true (default), the function only logs the actions that would be taken without actually performing them.
    Set to $false to perform the actual cleanup.

    .EXAMPLE
    Clear-PureStorageCloudFromAVS

    Runs in dry-run mode (default). Shows what would be removed without making any changes.

    .EXAMPLE
    Clear-PureStorageCloudFromAVS -DryRun $false

    Performs the actual cleanup, removing all Pure Storage Cloud resources from AVS.

    .INPUTS
    DryRun flag (boolean, defaults to $true)

    .OUTPUTS
    None
#>

function Clear-PureStorageCloudFromAVS {
    [CmdletBinding()]
    [AVSAttribute(60, UpdatesSDDC = $false, AutomationOnly = $false)]
    param(
        [Parameter(Mandatory = $false)]
        [bool]$DryRun = $true
    )

    if ($DryRun) {
        PrintLog "Running in DRY RUN mode - no changes will be made" WARNING
    } else {
        PrintLog "Running in LIVE mode - changes will be applied" WARNING
    }

    PrintLog "STEP 1: Cleaning up Datastores for vVols" INFO
    if ($DryRun) { Remove-VvolDatastore -DryRun } else { Remove-VvolDatastore }

    PrintLog "STEP 2: Cleaning up Datastores for VMFS" INFO
    if ($DryRun) { Remove-VmfsDatastore -DryRun } else { Remove-VmfsDatastore }

    PrintLog "STEP 3: Cleaning up Pure Storage iSCSI static targets" INFO
    if ($DryRun) { Remove-IscsiStaticTargets -DryRun } else { Remove-IscsiStaticTargets }

    PrintLog "STEP 4: Cleaning up Pure Storage plugin" INFO
    if ($DryRun) { Remove-PSRemotePlugin -DryRun } else { Remove-PSRemotePlugin }

    PrintLog "STEP 5: Cleaning up Pure Storage Storage Provider" INFO
    if ($DryRun) { Remove-PSStorageProvider -DryRun } else { Remove-PSStorageProvider }

    PrintLog "STEP 6: Cleaning up Pure Storage AVS service account" INFO
    if ($DryRun) { _Remove-AvsServiceAccount -Suffix "*" -DryRun } else { _Remove-AvsServiceAccount -Suffix "*" }

    PrintLog "STEP 7: Cleaning up Pure Storage storage policies" INFO
    if ($DryRun) { Remove-PSStoragePolicy -DryRun } else { Remove-PSStoragePolicy }

    PrintLog "All Pure Storage Cloud resources have been cleaned up from AVS" INFO
}