AzTablesWithManagedJwt.ps1

################################################################################
##                              Helper Functions                              ##
################################################################################

function DecodeJwt {
    param ($encodedToken)
    begin {
    filter From-Base64 {
            param ($b64)
            [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($b64 + ("=" * ($b64.Length % 4))))
        }
    }
    process {
        $header, $body, $sig = $encodedToken.Split('.')
        [pscustomobject]@{
            Header    = From-Base64 -b64 $header | ConvertFrom-Json
            Body      = From-Base64 -b64 $body| ConvertFrom-Json
            Signature = $sig
        }
    }
}


function ViewJwt {
    param (
        [string]
        $encodedToken
    )    
    $token = DecodeJwt -encodedToken $encodedToken
    
    Write-Host -ForegroundColor Yellow 'Header:'
    $token.Header | ConvertTo-Json

    Write-Host ""
    Write-Host -ForegroundColor Yellow 'Body:'
    $token.Body | ConvertTo-Json

    
    Write-Host ""
    Write-Host -ForegroundColor Yellow 'Signature:'
    $token.Signature
}


function DoWithRetry {
    param (
        [ScriptBlock]
        $Command,
        $RetryLimit=5,
        $Backoff=2,
        $ArgumentList
    )
    begin {
        function GetBackoffTime {
            # Binary exponential backoff 
            # if you've retried 100 times, do something else, this ain't happenin
            param ([ValidateRange(0,100)] $retries, $backoff=2)
            if ($retries -eq 0 -or $backoff -eq 0) { return 0 }
            [Math]::Pow($backoff, $retries)
        }
    }
    process {
        $retries   = 0
        $threshold = $RetryLimit
        $backoff   = $Backoff
        $tryAgain = $true

        :tryloop
        do {
            if ($threshold -le $retries) { break tryloop }
        
            $timeout = GetBackoffTime -retries $retries
            if ($timeout) { 
                Write-Debug "Waiting $timeout seconds after failure"
                Start-Sleep -Seconds $timeout
            }
        
            try {
                & $Command @ArgumentList
                                          
                $tryAgain = $false
            }
            catch {
                Write-Error $_
                $retries++
            }
        } while ($tryAgain)
    }
}


################################################################################
##                               Token Handling                               ##
################################################################################

function Epoch { param($When=[DateTime]::UtcNow) ([DateTimeOffset]([DateTime]$When)).ToUnixTimeSeconds() }

New-Variable -Scope Script -Name 'tokenInfo' -Force
#$RunningContext = 'VM'
$RunningContext = 'AzFn'
function RefreshJwtToken {

    $resource = "https://storage.azure.com/"

    if ($RunningContext -like 'VM') {
        $uri = "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=$resource"
        $response = $null

        $tokenCommand = {
            $response = Invoke-WebRequest -UseBasicParsing -Uri $uri -Headers @{Metadata="true"}
            $responseContent = $response.Content | ConvertFrom-Json
            $script:tokenInfo = $responseContent
        }
        DoWithRetry -Command $tokenCommand -ArgumentList @('2', 'baz') -RetryLimit 2
    }
    if ($RunningContext -like 'AzFn') {
        $script:tokenInfo = [pscustomobject]@{
            access_token=''
            token_type=''
            expires_on=''
            content=$null
        }

        $tokenCommand = {
            $script:tokenInfo.content = Get-AzAccessToken -Resource $resource
            $script:tokenInfo.access_token = $tokenInfo.content.Token
            $script:tokenInfo.token_type = $tokenInfo.content.Type
            $script:tokenInfo.expires_on = $tokenInfo.content.ExpiresOn.ToUnixTimeSeconds()
        }
        DoWithRetry -Command $tokenCommand -ArgumentList @('2', 'baz') -RetryLimit 2
    }
}

function IsTokenExpired {
    if (-not $tokenInfo) {
        throw "`$tokenInfo is null, Call RefreshJwtToken first."
    }
    # Check if it's good for more than the next ten seconds
    return $tokenInfo.expires_on -lt ((Epoch) - 10)
}

function RefreshIfExpired {
    # Get the token if you don't have one
    if (-not $tokenInfo) {
        RefreshJwtToken
    }
    # Refresh it if it's too old
    if (IsTokenExpired) {
        RefreshJwtToken
    }
}

function RefreshTokenIfExpired {
    function IsTokenExpired {
        if (-not $tokenInfo) {
            throw "`$tokenInfo is null, Call RefreshJwtToken first."
        }
        # Check if it's good for more than the next ten seconds
        return $tokenInfo.expires_on -lt ((Epoch) - 10)
    }

    # Get the token if you don't have one
    if (-not $tokenInfo) {
        RefreshJwtToken
    }
    # Refresh it if it's too old
    if (IsTokenExpired) {
        RefreshJwtToken
    }
}

################################################################################
##                              Az Table Handler                              ##
################################################################################



function MakeTableRequest {
    param (
        [string]
        $Account=$storageAccount,
        [string]
        $Table=$table,
        [string]
        $Query,
        $PartitionKey,
        $RowKey,
        $Top=-1,
        [ValidateSet('GET','POST','PUT','DELETE')]
        $Method='GET',
        [ValidateSet('None','Minimal','Full')]
        $OdataMetadata='None'
    )
    begin {
        ## Get access token if we don't have one or refresh if it expires in the next 10 seconds
        RefreshTokenIfExpired

        function EscapeString {
            param ([string]$inputString)
            [String]::Concat((
                $inputString.ToCharArray() | foreach {
                    switch ($_) {
                        ' ' { '%20' }
                        '/' { '%2F' }
                        '?'	{ '%3F' }
                        ':'	{ '%3A' }
                        '@'	{ '%40' }
                        '&' { '%26' }
                        '='	{ '%3D' }
                        '+' { '%2B' }
                        ',' { '%2C' }
                        '$' { '%24' }
                        default { $_ }
                    }
                }
            ))
        }
    }
    process {

        #region tableCall
        $odataAccept = switch ($OdataMetadata) {
            'Full'    { ';odata=fullmetadata' }
            'Minimal' { ';odata=minimalmetadata' }
            default   { '' }
        }
        $requestHeader = @{
            Origin=$(& hostname)
            Date=(Get-Date -Format "ddd, dd MMM yyyy HH:mm:ss 'GMT'")
            Accept='application/json' + $odataAccept
            'x-ms-version'='2017-11-09'
            Authorization="$($tokenInfo.token_type) $($tokenInfo.access_token)"
        }

        $quoteStart = [regex]"^'"
        $quoteEnd = [regex]"'`$"

        $uri = "$Account/$Table"
        $queryUri = @()
        if ($PartitionKey -and $RowKey) { $uri += "(PartitionKey='$(EscapeString $PartitionKey)',RowKey='$(EscapeString $RowKey)')" }
        if ($Query) {
            
            # Replace single quotes on the ends of tokens with the url encoded version %27
            # then escape innner quotes with ''. https://learn.microsoft.com/en-us/rest/api/storageservices/querying-tables-and-entities#single-quote-
            $Query = ($Query.Split(' ') | foreach { $quoteEnd.Replace($quoteStart.Replace($_,'%27'),'%27').Replace("'","''") }) -join ' '
            
            # Escape specific characters found in the filter query. https://learn.microsoft.com/en-us/rest/api/storageservices/querying-tables-and-entities#query-string-encoding
            $queryUri += "`$filter={0}" -f (EscapeString $Query)
        }
        if ($Top -ge 0) {
            $queryUri += "`$top=$Top"
        }
        
        $queryUri += "&_={0}" -f (Epoch)
        $uri = "{0}?{1}" -f $uri, ($queryUri -join '&')
        
        $response = try {
            (Invoke-RestMethod -UseBasicParsing -Uri $uri -Method $Method -Headers $requestHeader -ErrorAction Stop)
        } catch [System.Net.WebException] {
            $_.Exception.Response
        }

        return $response
        #endregion tableCall
    }
}


$storageAccount = "https://not.your.table.usgovcloudapi.net"
$table = "empid"

$data = MakeTableRequest -Account $storageAccount -Table $table
$data