Wrapper on Invoke-WebRequest to handle Citrix Cloud authentication, including token expiry

CitrixAuth.psm1

<#
MIT License

Copyright 2021, Citrix Systems, Inc.

Permission is hereby granted, free of charge, to any person obtaining a copy of 
this software and associated documentation files (the "Software"), to deal in 
the Software without restriction, including without limitation the rights to 
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 
of the Software, and to permit persons to whom the Software is furnished to do 
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all 
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#>

<#
 .Synopsis
  Initialize the CitrixAuth module.

 .Description
  Initialize the CitrixAuth module. This function stores the supplied client 
  id, client secret and token URL (if one was supplied), authenticates the 
  client with Citrix Cloud and, optionally, returns the resulting bearer token.

  The `Initialize-CitrixAuth` cmdlet stores the bearer token for subsequent use 
  by `Invoke-CloudRequest`.

 .Parameter ClientId
  The client id, obtained from the API Access tab in Citrix Cloud Identity and 
  Access Management.

 .Parameter ClientSecret
  The client secret, obtained from the API Access tab in Citrix Cloud Identity 
  and Access Management.

 .Parameter TokenUrl
  The OAuth 2.0 token endpoint URL for the Citrix Cloud Trust Service. Defaults
  to the US geo endpoint, "https://api-us.cloud.com/cctrustoauth2/root/tokens/clients". 
  Other possible values are "https://api-eu.cloud.com/cctrustoauth2/root/tokens/clients"
  for the EU geo and "https://api-ap-s.cloud.com/cctrustoauth2/root/tokens/clients" 
  for Asia Pacific South.

 .Parameter ReturnToken
  If set, the OAuth 2.0 bearer token is returned.

 .Example
  PS> Initialize-CitrixAuth -ClientId "abcd…" -ClientSecret "1234…"

  Initialize CitrixAuth with a client id and secret. 

 .Example
  PS> Initialize-CitrixAuth -ClientId "abcd…" -ClientSecret "1234…" -ReturnToken
  token_type access_token                                                                                                              
  ---------- ------------                                                                                                              
  bearer     eyJhbGciOiJ…

  Initialize CitrixAuth with a client id and secret, returning the bearer token.

 .Example
  PS> Initialize-CitrixAuth "abcd…" "1234…" "https://api-eu.cloud.com/cctrustoauth2/root/tokens/clients"

  Initialize CitrixAuth with a client id and secret, using the Citrix Cloud EU 
  token endpoint.
#>
function Initialize-CitrixAuth {
    param (
        [Parameter(Mandatory)]
        [string]
        $clientId,
        
        [Parameter(Mandatory)]
        [string] 
        $clientSecret,

        [Parameter()]
        [string] 
        $tokenUrl = 'https://api-us.cloud.com/cctrustoauth2/root/tokens/clients',

        [Parameter()]
        [switch] 
        $returnToken
    )

    # Save the client credentials in the script scope
    Set-Variable -Name client -Scope Script -Value @{
        ClientId     = $clientId
        ClientSecret = $clientSecret
        TokenUrl     = $tokenUrl
    }

    $token = Get-CloudToken

    if ($returnToken) {
        return $token
    }
}

<#
 .Synopsis
  Call a Citrix Cloud API.

 .Description
  The `Invoke-CloudRequest` cmdlet sends HTTP and HTTPS requests to a Citrix 
  Cloud API endpoint, returning a JSON-encoded response. The cmdlet sets
  the HTTP Authorization header with the bearer token obtained by 
  `Initialize-CitrixAuth`. The cmdlet transparently handles token expiry; when 
  an API call fails due to an invalid (i.e. expired) token, `Invoke-CloudRequest`
  automatically obtains a new token and retries the request. An informational
  message is displayed when the token is renewed.

  Note - you MUST pass a client id and secret to `Initialize-CitrixAuth` before
  using this cmdlet.

  The cmdlet wraps `Invoke-WebRequest` - see that cmdlet's documentation for
  supported parameters.

 .Example
  PS> Initialize-CitrixAuth -ClientId "abcd…" -ClientSecret "1234…"
  PS> Invoke-CloudRequest "https://registry.citrixworkspacesapi.net/acme/resourcelocations"

  StatusCode        : 200
  StatusDescription : OK
  Content           : {"items":[{"id":"2f86ae84-5025-4d4e-9eb5-273eb7ec81c7","name":"My Resource                                       
                    Location","internalOnly":false,"timeZone":"GMT Standard Time","readOnly":false}]}
  RawContent        : HTTP/1.1 200 OK
                    Cache-Control: no-store, must-revalidate, no-cache, max-age=0, private
                    Pragma: no-cache
                    Server: ********
                    Server: ********
                    X-Cws-TransactionId: dea2d5ac-8255-4bfe-960b-46152e408761
                    Acce…
  Headers           : {[Cache-Control, System.String[]], [Pragma, System.String[]], [Server, System.String[]], [X-Cws-TransactionId, 
                    System.String[]]…}
  Images            : {}
  InputFields       : {}
  Links             : {}
  RawContentLength  : 156
  RelationLink      : {}

  Initialize CitrixAuth with a client id and secret, then get a customer's 
  resource locations.

 .Example
  PS> Invoke-CloudRequest "https://registry.citrixworkspacesapi.net/acme/resourcelocations" -InformationAction Continue
  Invalid token - getting a new one

  StatusCode        : 200
  StatusDescription : OK
  Content           : {"items"…

  Get a customer's resource location, setting the `-InformationAction` parameter
  to `Continue`. In this example, the stored bearer token has expired, so 
  `Invoke-CloudRequest` catches the exception, obtains a new token, and retries
  the request.
#>
function Invoke-CloudRequest {
    [CmdletBinding(DefaultParameterSetName = 'StandardMethod', HelpUri = 'https:#go.microsoft.com/fwlink/?LinkID=2097126')]
    param(
        [switch]
        ${UseBasicParsing},

        [Parameter(Mandatory = $true, Position = 0)]
        [ValidateNotNullOrEmpty()]
        [uri]
        ${Uri},

        [Microsoft.PowerShell.Commands.WebRequestSession]
        ${WebSession},

        [Alias('SV')]
        [string]
        ${SessionVariable},

        [switch]
        ${AllowUnencryptedAuthentication},

        [Microsoft.PowerShell.Commands.WebAuthenticationType]
        ${Authentication},

        [pscredential]
        [System.Management.Automation.CredentialAttribute()]
        ${Credential},

        [switch]
        ${UseDefaultCredentials},

        [ValidateNotNullOrEmpty()]
        [string]
        ${CertificateThumbprint},

        [ValidateNotNull()]
        [X509Certificate]
        ${Certificate},

        [switch]
        ${SkipCertificateCheck},

        [Microsoft.PowerShell.Commands.WebSslProtocol]
        ${SslProtocol},

        [securestring]
        ${Token},

        [string]
        ${UserAgent},

        [switch]
        ${DisableKeepAlive},

        [ValidateRange(0, 2147483647)]
        [int]
        ${TimeoutSec},

        [System.Collections.IDictionary]
        ${Headers},

        [ValidateRange(0, 2147483647)]
        [int]
        ${MaximumRedirection},

        [ValidateRange(0, 2147483647)]
        [int]
        ${MaximumRetryCount},

        [ValidateRange(1, 2147483647)]
        [int]
        ${RetryIntervalSec},

        [Parameter(ParameterSetName = 'StandardMethodNoProxy')]
        [Parameter(ParameterSetName = 'StandardMethod')]
        [Microsoft.PowerShell.Commands.WebRequestMethod]
        ${Method},

        [Parameter(ParameterSetName = 'CustomMethodNoProxy', Mandatory = $true)]
        [Parameter(ParameterSetName = 'CustomMethod', Mandatory = $true)]
        [Alias('CM')]
        [ValidateNotNullOrEmpty()]
        [string]
        ${CustomMethod},

        [Parameter(ParameterSetName = 'StandardMethodNoProxy', Mandatory = $true)]
        [Parameter(ParameterSetName = 'CustomMethodNoProxy', Mandatory = $true)]
        [switch]
        ${NoProxy},

        [Parameter(ParameterSetName = 'StandardMethod')]
        [Parameter(ParameterSetName = 'CustomMethod')]
        [uri]
        ${Proxy},

        [Parameter(ParameterSetName = 'StandardMethod')]
        [Parameter(ParameterSetName = 'CustomMethod')]
        [pscredential]
        [System.Management.Automation.CredentialAttribute()]
        ${ProxyCredential},

        [Parameter(ParameterSetName = 'StandardMethod')]
        [Parameter(ParameterSetName = 'CustomMethod')]
        [switch]
        ${ProxyUseDefaultCredentials},

        [Parameter(ValueFromPipeline = $true)]
        [System.Object]
        ${Body},

        [System.Collections.IDictionary]
        ${Form},

        [string]
        ${ContentType},

        [ValidateSet('chunked', 'compress', 'deflate', 'gzip', 'identity')]
        [string]
        ${TransferEncoding},

        [string]
        ${InFile},

        [string]
        ${OutFile},

        [switch]
        ${PassThru},

        [switch]
        ${Resume},

        [switch]
        ${SkipHttpErrorCheck},

        [switch]
        ${PreserveAuthorizationOnRedirect},

        [switch]
        ${SkipHeaderValidation},

        [switch]
        ${Retry})

    begin {
        try {
            # Boilerplate generated code
            $outBuffer = $null
            if ($PSBoundParameters.TryGetValue('OutBuffer', [ref]$outBuffer)) {
                $PSBoundParameters['OutBuffer'] = 1
            }

            # Wrap Invoke-WebRequest
            $wrappedCmd = $ExecutionContext.InvokeCommand.GetCommand('Microsoft.PowerShell.Utility\Invoke-WebRequest', [System.Management.Automation.CommandTypes]::Cmdlet)

            # Get the Headers parameter, if one has been passed, or an empty 
            # hash table.
            if ($PSBoundParameters.TryGetValue('Headers', [ref]$headers)) {
                # Remove the headers parameter, since we'll pass a new one
                $PSBoundParameters.Remove('Headers') | Out-Null
            }
            else {
                $headers = @{}
            }

            # Get the stored authorization header and put it in the headers
            $auth = Get-Variable -Name authorization -Scope Script -ValueOnly
            if ($auth) {
                $headers["Authorization"] = $auth
            }
            else {
                throw "You must call Initialize-CitrixAuth with the client id and secret before calling Invoke-CloudRequest"
            }

            # Grab the Retry parameter, and don't pass it to Invoke-WebRequest
            if ($PSBoundParameters.TryGetValue('Retry', [ref]$retry)) {
                $PSBoundParameters.Remove('Retry') | Out-Null
            }
            else {
                $retry = $FALSE                
            }

            # Build a script block with our new headers and the remaining parameters
            $scriptCmd = { & $wrappedCmd -Headers $headers @PSBoundParameters }

            # Start the cmdlet pipeline
            $steppablePipeline = $scriptCmd.GetSteppablePipeline($myInvocation.CommandOrigin)
            $steppablePipeline.Begin($PSCmdlet)
        }
        catch {
            throw
        }
    }

    process {
        try {
            $steppablePipeline.Process($_)
        }
        catch [Microsoft.PowerShell.Commands.HttpResponseException] {
            # Only retry once, and only for the 401 Unauthorized status code
            if (-not $retry -and $_.Exception.Response.StatusCode -eq [System.Net.HttpStatusCode]::Unauthorized) {
                Write-Information "Invalid token - getting a new one"

                # We don't need the returned token
                Get-CloudToken | Out-Null

                # Set retry so that we don't infinitely recurse
                Invoke-CloudRequest @PSBoundParameters -Retry

                $recursed = $TRUE
            }
            else {
                throw
            }
        }
        catch {   
            throw
        }
    }

    end {
        if ($recursed) {
            # Don't end the pipeline, since that will report the error that we caught!
            $steppablePipeline.Dispose()
        }
        else {
            try {
                $steppablePipeline.End()
            }
            catch {
                throw
            }        
        }
    }
}

# Internal function to use stored client credentials to get a new bearer token
function Get-CloudToken() {
    # Get the stored client credentials
    $client = Get-Variable -Name client -Scope Script -ValueOnly

    try {
        # OAuth 2.0 authorization
        $response = Microsoft.PowerShell.Utility\Invoke-WebRequest $client['TokenUrl'] -Method POST -Body @{
            grant_type    = "client_credentials"
            client_id     = $client['ClientId']
            client_secret = $client['ClientSecret']
        }
    }
    catch [Microsoft.PowerShell.Commands.HttpResponseException] {
        # Don't try to parse $response if there was an error
        throw
    }

    $token = $response.Content | ConvertFrom-Json

    # Store the HTTP Authorization header for subsequent use by Invoke-CloudRequest
    Set-Variable -Name authorization -Scope Script -Value "CwsAuth Bearer=$($token.access_token)"

    return $token
}

<#
 .Synopsis
  Test invalidation of the stored bearer token.

 .Description
  This function sets the stored bearer token to the string "DUMMY", simulating
  an expired token. The next call of `Invoke-CloudRequest` will get a new bearer
  token.

 .Example
  PS> Test-CloudTokenExpiry
  PS> Invoke-CloudRequest "https://registry.citrixworkspacesapi.net/acme/resourcelocations" -InformationAction Continue
  Invalid token - getting a new one

  StatusCode        : 200
  StatusDescription : OK
  Content           : {"items"…

  Invalidate the stored bearer token, then set `-InformationAction Continue` 
  so that the token refresh is visible.
#>
function Test-CloudTokenExpiry() {
    Set-Variable -Name authorization -Scope Script -Value "CwsAuth Bearer=DUMMY"
}

Export-ModuleMember -Function Initialize-CitrixAuth
Export-ModuleMember -Function Invoke-CloudRequest
Export-ModuleMember -Function Test-CloudTokenExpiry