GraphFastVoleer.psm1

<#
.Synopsis
   Executes Microsoft Graph queries quickly
.DESCRIPTION
   This cmdlet allows you to run bulk Microsoft Graph API requests quickly. Internally it uses both multi-threading (via PS runspaces) and implements the Odata batching functionality of the Graph API. This can result in speed increases of up to and over 60x.
.EXAMPLE
   Invoke-GraphFast -ClientId 12345678-1234-1234-1234-123456789012 -ClientSecret "324trgh7b4b3a!sgfn3p9757a9ewhg7a" -TenantId 12345678-1234-1234-1234-123456789012 -Urls ("/teams/[TEAMID]","/teams/[TEAMID]","/teams/[TEAMID]",.....)
#>

function Invoke-GraphFast {
    [CmdletBinding(DefaultParameterSetName = 'TokenAuth')]
    param(
        # A Microsoft Graph access token. If specified with a ClientID, the cmdlet will derive the tenantId from the token, even if the TeantID parameter is provided
        [Parameter(ParameterSetName = 'TokenAuth', Mandatory = $true)] $AccessToken,
        # A Microsoft Graph refresh token
        [Parameter(ParameterSetName = 'TokenAuth', Mandatory = $false)] $RefreshToken,
        # The ClientId of a Microsoft Azure AD Application Registration
        [Parameter(ParameterSetName = 'CertAuth', Mandatory = $true)][Parameter(ParameterSetName = 'CertPath', Mandatory = $true)][Parameter(ParameterSetName = 'ClientAuth', Mandatory = $true)] $ClientId,
        # An authentication secret of the Microsoft Azure AD Application Registration
        [Parameter(ParameterSetName = 'ClientAuth', Mandatory = $true)] $ClientSecret,
        # An authentication certificate linked to the Microsoft Azure AD Application Registration
        [Parameter(ParameterSetName = 'CertAuth', Mandatory = $true)][System.Security.Cryptography.X509Certificates.X509Certificate2] $ClientCertificate,
        # The Microsoft Azure AD TenantId (GUID or domain)
        [Parameter(ParameterSetName = 'CertAuth', Mandatory = $true)][Parameter(ParameterSetName = 'CertPath', Mandatory = $true)][Parameter(ParameterSetName = 'ClientAuth', Mandatory = $true)]  $TenantId,
        # The path to a certificate PFX file that has it's corresponding .CER file linked to the Microsoft Azure AD Application Registration
        [Parameter(ParameterSetName = 'CertPath', Mandatory = $true)] $CertificatePath,
        # The certificate password for the PFX file provided in the CertPath parameter
        [Parameter(ParameterSetName = 'CertPath', Mandatory = $true)] $CertificatePassword,
        # The Microsoft Graph endpoint to use (v1.0 or beta)
        $Endpoint = "v1.0",
        # A single Graph URL or a collection of URLs. This may be the full Graph URL or may begin with the portion after the endpoint (eg. /users/myuser@mydomain.com). If the full URL is provided, the endpoint is ignored and the value of the -Endpoint parmeter is used.
        $Urls,
        # Sets the ConsistencyLevel header for Graph calls that need it.
        $ConsistencyLevel,
        # Switch to disable retrieving paged results
        [Switch] $NoPaging,
        # Sets the PowerShell runspace pool size. By default this is one more than the number of logical processors
        $PoolSize = (1 + $Env:NUMBER_OF_PROCESSORS),
        # The maximum number of times to retry requests that return any status code other than 200, 403, 404, or 429
        $MaxRetries = 3,
        # If specified, the raw results of the batch response will be returned with all requests, including ones that failed
        $ReturnErrors = $false, #If set to true, this will return the content of non 200 responses
        # If specified, the raw results of the batch response will be returned with all requests that returned status code 200
        $ReturnAll = $false, #If set to false, this will return the full response and not just the body,
        # By default, if the number of URLs provided is over 100, this cmdlet will output status updates on the Information stream
        $ShowProgress = $true,
        # When ShowProgress is true, if this value is set, update messages will only be displayed every X minutes
        $UpdateInterval,
        # When ShowProgress is true, if this value is not set to false, at the end, a total objects processed message will be displayed.
        $ShowTotals = $true,
        # When ShowProgress is true, this value will be used to populate the output messages eg. "Total groups processed: 100 of 2134"
        $ItemType = "objects",
        # Used internally to pass a collection of GraphBatch objects instead of URLs.
        $GraphBatch,
        # Used internally to track retrys
        [Parameter(DontShow)] $retry = 0, # This should not be set by users, it will be set by recursive calls
        # Used internally to alter functionality when paging results
        [Parameter(DontShow)] $large = $false,
        # Used internally to handle credentials during recursive calls
        [Parameter(DontShow)] $tokens
    )
    Begin {
       
        function Test-AccessToken ($tokens) {
            $tokens.jwt = Get-JWTDetails $tokens.accessToken -Verbose:$false
            if ($tokens.jwt.TimeToExpiry -lt (New-TimeSpan -Minutes 5)) {
                if ([String]::IsNullOrEmpty($tokens.refreshToken) -eq $false) {
                    #If there's a refresh token, update the tokens
                    $tokens = Update-Tokens $tokens
                }
                elseif ($tokens.jwt.TimeToExpiry -gt 0) {
                    Write-Warning "Access token is about to expire in $($tokens.jwt.TimeToExpiry.minutes) minutes and cannot be refreshed."
                }
                else {    
                    throw "Access token expired. No refresh token available to update the access token." 
                }
                
            }
            Write-Output $tokens
            $PSDefaultParameterValues."Invoke-RestMethod:Headers" = @{Authorization = "Bearer $($tokens.accessToken)" }
        }

        function Update-Tokens ($tokens) {
            $body = @{
                client_id     = $tokens.jwt.appid
                scope         = $tokens.jwt.scp
                grant_type    = "refresh_token"
                refresh_token = $tokens.refreshToken   
            }
            $return = Invoke-RestMethod "https://login.microsoftonline.com/organizations/oauth2/v2.0/token" -Method POST -Body $body
            $tokens.RefreshToken = $return.refresh_token
            $tokens.AccessToken = $return.access_token
            $tokens.jwt = Get-JWTDetails $tokens.accessToken -Verbose:$false
            $PSDefaultParameterValues."Invoke-RestMethod:Headers" = @{Authorization = "Bearer $($tokens.accessToken)" }
            $tokens
        }   
    }

    process {


        function Write-Host ($m) { 
            $context.SaveMessage("Information", $m) 
            Microsoft.PowerShell.Utility\Write-Host $m
        }
        function Write-Information ($m) { 
            $context.SaveMessage("Information", $m) 
            Microsoft.PowerShell.Utility\Write-Information $m
        }
        function Write-Warning ($m) { 
            $context.SaveMessage("Information", $m) 
            Microsoft.PowerShell.Utility\Write-Warning $m
        }
        trap {
            Write-Error "UNHANDLED EXCEPTION" 
            Write-Error "----------------------------------------"
            Write-Error "Error: $($_ | Out-String)"
        }
        
        if ($null -eq $endpoint) { $endpoint = "v1.0" }

        #Check for Access Token and validity and get a new token if needed
        if ($PSCmdlet.ParameterSetName -eq 'CertPath') {
            if ($CertificatePassword.GetType() -eq [String]) { $CertificatePassword = ConvertTo-SecureString $CertificatePassword -AsPlainText -Force }
            if ($CertificatePath -match "^voleer://") { $CertificatePath = $context.DownloadFile($CertificatePath) }
            $clientCertificate = New-Object System.Security.Cryptography.X509Certificates.X509Certificate2
            $clientCertificate.import($CertificatePath, $CertificatePassword, [System.Security.Cryptography.X509Certificates.X509KeyStorageFlags]::DefaultKeySet)
        }
        
        switch ($PSCmdlet.ParameterSetName) {
            'ClientAuth' {
                if ($clientSecret.GetType() -eq [String]) { $clientSecret = ConvertTo-SecureString $clientSecret -AsPlainText -Force }
                $clientApplication = New-MsalClientApplication -ClientId $clientId -ClientSecret $clientSecret -TenantId $tenantId
            }
            'CertAuth' {
                $clientApplication = New-MsalClientApplication -ClientId $clientId -ClientCertificate $clientCertificate -TenantId $tenantId
            }
            'CertPath' {
                $clientApplication = New-MsalClientApplication -ClientId $clientId -ClientCertificate $clientCertificate -TenantId $tenantId
            }
            'TokenAuth' {
                #Populate Tokens Object if it isn't already
                if ($null -eq $tokens) {
                    $tokens = [PSCustomObject]@{
                        AccessToken  = $AccessToken
                        refreshToken = $refreshToken
                        jwt          = try { Get-JWTDetails $AccessToken -Verbose:$false } catch { $null };
                    }
                }
                $tokens = Test-AccessToken -tokens $tokens
            }
        }


        #Populate GraphBatch with urls, ignoring duplicates
        if ($null -eq $GraphBatch) {
            $GraphBatch = [System.Collections.ArrayList]@()
        }
        else {
            $newBatch = [System.Collections.ArrayList]@()
            foreach ($g in $graphBatch) {
                if ($g.GetType().name -eq "GraphBatch") {
                    if ($g.Method.ToUpper() -in ("POST", "PUT", "PATCH", "MERGE")) {
                        [void]$newBatch.Add(($g | Select-Object *))
                    }
                    elseif ($g.Method.ToUpper() -in ("DEFAULT", "DELETE", "GET", "HEAD", "OPTIONS")) {
                        [void]$newBatch.Add(($g | Select-Object id, method, url))
                    }
                    else {
                        throw "Invalid method specified in GraphBatch object: $($g.method)"
                    }
                }
                else {
                    #Custom Logic to allow non classed objects with valid fields through
                }
            }
            $graphBatch = $newBatch
        }
        
        if ($null -ne $urls) {
            $urlHashTable = @{}
            $idHashTable = @{}
            #Add any GraphBatch objects to the hashtable to prevent incoming urls from creating duplicates
            foreach ($g in $graphBatch) { $urlHashTable[$g.url] = $null; $idHashTable[$g.id] = $null }
                
            foreach ($url in $Urls) {
                try {
                    $urlHashTable.add($url, $null)
                    try { $idHashTable.add($url, $null) } catch { Write-Warning "Duplicate id in the GraphBatch collection provided."; throw }
                    if ($ConsistencyLevel) {
                        [void] $GraphBatch.Add([GraphBatch]::new($url, "GET", $url, @{ConsistencyLevel = "Eventual" }))
                    }
                    else {
                        [void] $GraphBatch.Add([GraphBatch]::new($url))    
                    }
                    
                }
                catch {
                    Write-Warning "Unable to add duplicate url: $url"
                }
            }
        }
        
       
        #Create RunspacePool
        $runspacePool = [runspacefactory]::CreateRunspacePool(1, $poolSize)
        $runspacePool.Open()        
        $jobs = [System.Collections.ArrayList]@()
        
        $contentType = "application/json"
        $uri = "https://graph.microsoft.com/$endpoint/`$batch"
    
        #Start Running Batches
        $x = 0
        $runningCount = 0
        $pause = 1100
        $stopInterval = 100
        if ($UpdateInterval) {
            $UpdateInterval = New-TimeSpan -Minutes $UpdateInterval
            $timer = New-Object -TypeName System.Diagnostics.Stopwatch
            $timer.Start()
        }
        if ($ConsistencyLevel) { try { $PSDefaultParameterValues."Invoke-RestMethod:Headers".Add("ConsistencyLevel", $ConsistencyLevel) } catch {} }
        for ($x = 0; $x -lt $graphBatch.count; $x += 20) {
            if ($showProgress) { $runningCount += $graphBatch[$x..($x + 19)].count }

            $body = [PSCustomObject] @{requests = $graphBatch[$x..($x + 19)] } | ConvertTo-Json -Depth 10
            $body = $body.Replace("\u0027", "'").replace("\u0026", "&")
            $instance = [powershell]::Create().AddScript( { param($header, $uri, $body, $contentType); Invoke-RestMethod -Method POST -Headers $header -Uri $uri -Body $body -ContentType $contentType })
            if ($tokens.accessToken) { [void] ($instance.AddArgument($PSDefaultParameterValues."Invoke-RestMethod:Headers")) }
            else { [void] ($instance.AddArgument(@{Authorization = ($clientApplication | Get-MSALToken).CreateAuthorizationHeader() })) }
            [void] ($instance.AddArgument($uri))
            [void] ($instance.AddArgument($body))
            [void] ($instance.AddArgument($contentType))
            [void] ($instance.RunspacePool = $runspacePool)
            [void] $jobs.Add( [PSCustomObject]@{
                    instance = $instance;
                    job      = $instance.BeginInvoke()
                }
            )
            Start-Sleep -milliseconds 360   #Near optimal dely to prevent throttling

            if (($x + 20) % $stopInterval -eq 0 -and $x -ne 0) {
                if ($showProgress) {
                    if ($UpdateInterval) {
                        if ($timer.Elapsed -gt $UpdateInterval) {
                            Write-Information "Total $itemType processed: $runningCount of $($graphBatch.Count)" 
                            $timer.Restart()
                        }
                    }
                    else {
                        Write-Information "Total $itemType processed: $runningCount of $($graphBatch.Count)" 
                    }    
                }
                

                do {
                    $pendingjobs = $jobs | Where-Object { $_.job.isCompleted -eq $false }
                    if (@($pendingjobs).count -ne 0) {
                        #Write-Information "Waiting for $(@($pendingjobs).count) jobs"
                        Start-Sleep -Milliseconds 500
                    }
                }
                while ($false -in $jobs.job.IsCompleted)
                start-sleep -milliseconds $pause

                $currentresponses = [System.Collections.ArrayList]@()
                foreach ($job in $jobs) {
                    foreach ($response in $job.instance.EndInvoke($job.job).responses) {
                        [void] $currentresponses.add($response)
                    }
                }

                $throttledResponses = $currentresponses | Select-Object -last 100 | Where-Object status -eq "429"
                if ($throttledResponses) {
                    
                    if ($null -ne $throttledResponses.headers."retry-after") {
                        $recommendedWait = ($throttledResponses.headers | Measure-object "retry-after" -Maximum -ErrorAction "SilentlyContinue").maximum
                    }
                    if ($null -eq $recommendedWait) { 
                        #Added default just in case a retry-after header isn't present (rare edge case)
                        $recommendedWait = 30 
                    }
                    #Write-Information "Sleeping $recommendedWait seconds for too many requests. Request Count: $($throttledResponses.count)"
                    Start-Sleep -Seconds ($recommendedWait + ($pause / 1000))
                    $pause += 200


                }
                else {
                    if ($pause -gt 1100) { $pause -= 10 }

                }
                if ($tokens) { $tokens = Test-AccessToken -tokens $tokens }
            }
        }
        
        #Wait for Jobs to Complete
        while ($false -in $jobs.job.IsCompleted) { Start-sleep -Milliseconds 500 }
        
        $responses = [System.Collections.ArrayList]@()
        foreach ($job in $jobs) {
            foreach ($response in $job.instance.EndInvoke($job.job).responses) {
                [void] $responses.add($response)
            }
        }
        $runspacePool.close()
        $runspacePool = $null
        
        $tooManyRequests = 0
        $retries = [System.Collections.ArrayList]@()
        foreach ($response in $responses) {
            switch ($response.status) {
                200 { }
                403 { }
                404 { }
                429 {
                    [void] $retries.Add([GraphBatch]::new($response.id))
                    $tooManyRequests++
                }
                default {
                    [void] $retries.Add([GraphBatch]::new($response.id))
                }
            }
        }
        if ($retries.count -gt 0 -and $retry -lt $maxretries) {
            if ($tooManyRequests -gt 0) {
                Write-Information "Sleeping for too many requests. Count: $($tooManyRequests)"

            }
            Write-Verbose "Retrying $($retries.count) queries:`n$($responses | Group-Object status | out-string)"
            if ($VerbosePreference -eq "Continue") { $verbose = $true } else { $verbose = $false }
            $retryValue = $retry
            if ($tooManyRequests -eq 0) {
                $retryValue++
            }
            $params = @{
                GraphBatch       = $retries 
                maxRetries       = $maxretries 
                Endpoint         = $endpoint
                ReturnAll        = $true 
                showProgress     = $false 
                ConsistencyLevel = "$ConsistencyLevel"
                Retry            = $retryValue 
            }
            switch ($PSCmdlet.ParameterSetName) {
                { $_ -eq 'CertAuth' -or $_ -eq 'CertPath' } {
                    $params.add("ClientId", $clientId)
                    $params.add("ClientCertificate", $clientCertificate)
                    $params.add("TenantId", $tenantId)
                }
                'ClientAuth' {
                    $params.add("ClientId", $clientId)
                    $params.add("clientSecret", $clientSecret)
                    $params.add("TenantId", $tenantId)
                }
                'TokenAuth' {
                    $params.Add("AccessToken", $tokens.AccessToken)
                    $params.Add("RefreshToken", $tokens.RefreshToken)
                    $params.Add("tokens", $tokens)
                }
            }
            
            $retryresponses = Invoke-GraphFast @params -Verbose:$Verbose
            
        }

        $results = @{}
        foreach ($r in $responses) { $results.add($r.id, $r) }
        $successes = $retryResponses | Group-Object status | Where-Object name -eq 200
        foreach ($success in $successes.group) {
            $results[$success.id] = $success
        }

        if (-not $NoPaging) {
            #Now get all responses that have @odata.nextlinks and append the data until there's nothing left
            while ( $results.values | Where-Object { $null -ne $_.body."@odata.nextLink" -and $large -eq $false }) {
                $incomplete = [System.Collections.ArrayList]@()
                foreach ($incompleteBody in $results.values | Where-Object { $null -ne $_.body."@odata.nextLink" }) {
                    [void] $incomplete.add([GraphBatch]::new($incompleteBody.id, $incompleteBody.body."@odata.nextLink"))
                }
                if ($VerbosePreference -eq "Continue") { $verbose = $true } else { $verbose = $false }
                #Write-Host "." -nonewline
                $params = @{
                    GraphBatch       = $incomplete 
                    maxRetries       = $maxretries 
                    Endpoint         = $endpoint
                    ReturnAll        = $true 
                    showProgress     = $false 
                    ConsistencyLevel = "$ConsistencyLevel"
                    Retry            = $retry 
                    large            = $true
                }
                switch ($PSCmdlet.ParameterSetName) {
                    { $_ -eq 'CertAuth' -or $_ -eq 'CertPath' } {
                        $params.add("ClientId", $clientId)
                        $params.add("ClientCertificate", $clientCertificate)
                        $params.add("TenantId", $tenantId)
                    }
                    'ClientAuth' {
                        $params.add("ClientId", $clientId)
                        $params.add("clientSecret", $clientSecret)
                        $params.add("TenantId", $tenantId)
                    }
                    'TokenAuth' {
                        $params.Add("AccessToken", $tokens.AccessToken)
                        $params.Add("RefreshToken", $tokens.RefreshToken)
                        $params.Add("tokens", $tokens)
                    }
                }
                $largeResponses = Invoke-GraphFast @params -Verbose:$Verbose

                foreach ($response in $largeResponses) {
                    $results[$response.id].body.value += $response.body.value
                    try {
                        $results[$response.id].body."@odata.nextLink" = $null
                        $results[$response.id].body."@odata.nextLink" = $response.body."@odata.nextLink"
                    }
                    catch {}
                }
            }
        }
        
        
        if ($showProgress -eq $true -and $large -eq $false -and $showTotals -eq $true) { 
            $ValueNodeChildren = (($results.values | Where-Object { $_.status -eq 200 -and $null -ne $_.body.value }).body.value | Measure-Object).Count
            $queriesWithOutValueNode = ($results.values | Where-Object { $_.status -eq 200 -and $null -eq $_.body.value } | Measure-Object).Count
            Write-Information "Total $itemType returned: $($queriesWithOutValueNode + $ValueNodeChildren)"
        }
        if ($returnErrors) {
            if ($returnAll -eq $false) {
                Write-Output $results.values.body
            }
            else {
                Write-Output $results.values    
            }
            
        }
        else {
            if ($returnAll -eq $false) {
                Write-Output ($results.values | Where-Object status -eq 200).body
            }
            else {
                Write-Output ($results.values | Where-Object status -eq 200)
            }
            
        }
    }
    End {
        if ($runspacePool) { [void]($runspacePool.Close()) }
    }
}

function New-GraphBatchItem {
    param(
        [Parameter(Mandatory = $true)] $url,
        [Parameter(Mandatory = $false)]$id ,
        [Parameter(Mandatory = $false)]$method,
        [Parameter(Mandatory = $false)]$body,
        [Parameter(Mandatory = $false)]$headers
    )

    if ($null -ne $body) {
        [GraphBatch]::new($id, $method, $url, $body, $headers)
    }
    elseif ($null -ne $headers) {
        [GraphBatch]::new($id, $method, $url, $headers)
    }
    elseif ($null -ne $method) {
        [GraphBatch]::new($id, $method, $url, $headers)
    }
    elseif ($null -ne $id) {
        [GraphBatch]::new($id, $url)
    }
    else {
        [GraphBatch]::new($url)    
    }
}

Export-ModuleMember "Invoke-GraphFast"
Export-ModuleMember "New-GraphBatchItem"