Enhance token management and add Disconnect-AzResourceGraph function (#5)

* Enhance token management and add Disconnect-AzResourceGraph function

* Add tests and update changelog

Author: SimonWahlin

AzResourceGraph/Private/Assert-AzureConnection.ps1

@@ -65,6 +65,7 @@ function Assert-AzureConnection {
         # We can therefore safely remove the Interactive parameter from the local LocalTokenSplat
         if ($LocalTokenSplat.ContainsKey('Interactive')) {
             $LocalTokenSplat.Remove('Interactive')
+            $LocalTokenSplat.Remove('ClientId')
         }
 
         $NewToken = Get-AzToken @LocalTokenSplat -ErrorAction 'Stop'

AzResourceGraph/Private/Test-AzureToken.ps1

@@ -46,7 +46,12 @@ function Test-AzureToken {
     )
     return (
         $null -ne $Token -and
-        $Token.ExpiresOn -ge [System.DateTimeOffset]::Now.AddMinutes($MinValid) -and
-        $Token.Claims['aud'] -eq $Resource
+        $Token.ExpiresOn.UtcDateTime -ge [System.DateTimeOffset]::Now.AddMinutes($MinValid).UtcDateTime -and
+        (
+            # Due to AzAuth not showing correct aud claim when using Interactive and TokenCache we check either aud or scope
+            # https://github.com/PalmEmanuel/AzAuth/issues/137
+            $Resource -eq $Token.Claims['aud'] -or
+            $Token.Scopes -contains "$Resource/.default"
+        )
     )
 }

AzResourceGraph/Public/Connect-AzResourceGraph.ps1

@@ -32,6 +32,11 @@ Client secret string associated with the service principal.
 Switch indicating that the command should acquire a token using the
 Azure Managed Identity assigned to the current VM / App Service / container.
 
+.PARAMETER ManagementEndpoint
+Endpoint used for management. This is used for the Audience claim when authenticating to Azure.
+For global Azure, this should be left as the default of 'https://management.azure.com'.
+For Azure China, use 'https://management.chinacloudapi.cn' and for US Government Cloud use 'https://management.usgovcloudapi.net'.
+
 .EXAMPLE
 # 1. Interactive sign-in (prompts user)
 Connect-AzResourceGraph
@@ -82,13 +87,20 @@ function Connect-AzResourceGraph {
         [string]$ClientSecret,
 
         [Parameter(Mandatory, ParameterSetName = 'ManagedIdentity')]
-        [switch]$ManagedIdentity
+        [switch]$ManagedIdentity,
+
+        [Parameter(ParameterSetName = 'ManagedIdentity')]
+        [Parameter(ParameterSetName = 'Interactive')]
+        [Parameter(ParameterSetName = 'Certificate')]
+        [Parameter(ParameterSetName = 'ClientSecret')]
+        $ManagementEndpoint = 'https://management.azure.com'
     )
 
     # Set up module-scoped variables for getting tokens
     $script:TokenSplat = @{}
     $script:CertificatePath = $null
 
+    $script:TokenSplat['Resource'] = $ManagementEndpoint
     $script:TokenSplat['ClientId'] = $ClientId
     if ($PSBoundParameters.ContainsKey('Tenant')) {
         $script:TokenSplat['TenantId'] = $Tenant
@@ -109,6 +121,7 @@ function Connect-AzResourceGraph {
     }
     if ($PSCmdlet.ParameterSetName -eq 'Interactive') {
         $script:TokenSplat['Interactive'] = $true
+        $script:TokenSplat['TokenCache'] = 'AzResourceGraph'
     }
     if ($ManagedIdentity.IsPresent) {
         $script:TokenSplat['ManagedIdentity'] = $true
@@ -117,4 +130,7 @@ function Connect-AzResourceGraph {
     $script:Token = Get-AzToken @script:TokenSplat
     # Save the source of the token to module scope for AssertAzureConnection to know how to refresh it
     $script:TokenSource = 'Module'
+    if ($script:TokenSplat['Interactive'] -eq $true) {
+        $script:TokenSplat['UserName'] = $script:Token.Identity
+    }
 }

AzResourceGraph/Public/Disconnect-AzResourceGraph.ps1

@@ -0,0 +1,21 @@
+<#
+.SYNOPSIS
+Disconnects the current PowerShell session from Azure Resource Graph and clears cached access token in memory.
+
+.DESCRIPTION
+Disconnect-AzResourceGraph removes any stored session information and resets sessions configuration.
+
+.EXAMPLE
+# 1. Disconnect and clear session information
+Disconnect-AzResourceGraph
+
+#>
+function Disconnect-AzResourceGraph {
+    [CmdletBinding()]
+    param ()
+
+    $script:TokenSplat = @{}
+    $script:TokenSource = 'Global'
+    $script:Token = $null
+    $script:CertificatePath = $null
+}

CHANGELOG.md

@@ -5,8 +5,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### Added
-- Initial release of AzResourceGraph module.
+### Fixed
+- Fixed Auth problems
+
+## [0.1.1] - 2025-05-14
 
 ### Changed
 - Removed space from tags in module manifest
+
+## [0.1.0] - 2025-05-14
+
+### Added
+- Initial release of AzResourceGraph module.

tests/Unit/Public/Disconnect-AzResourceGraph.tests.ps1

@@ -0,0 +1,30 @@
+BeforeAll {
+    Import-Module -FullyQualifiedName "$PSScriptRoot/../../../output/module/AzResourceGraph"
+}
+
+Describe 'Disconnect-AzResourceGraph' {
+    InModuleScope 'AzResourceGraph' {
+        BeforeEach {
+            # Mock the Get-AzToken function to avoid actual authentication
+            Mock 'Get-AzToken' {
+                @{
+                    Token = 'dummy'
+                    ExpiresOn = [System.DateTimeOffset]::Now.AddHours(1)
+                    Claims = @{
+                        aud = 'myApp'
+                    }
+                }
+            }
+        }
+        Context 'Successful token removal' {
+            It 'Resets any token information' {
+                Connect-AzResourceGraph
+                Disconnect-AzResourceGraph
+                $script:TokenSplat.Keys.Count | Should -be 0
+                $null -eq $script:CertificatePath | Should -be $true
+                $null -eq $script:Token | Should -be $true
+            }
+        }
+
+    }
+}