From 5c0cafbb60eb2d0c1ac8c4e57859129b189bc063 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20W=C3=A5hlin?= <simon.wahlin@unwonted.se>
Date: Thu, 3 Jul 2025 12:42:08 +0200
Subject: [PATCH] Add Token parameter to Search-AzResourceGraph command

---
 .../Public/Search-AzResourceGraph.ps1            | 16 +++++++++++++---
 CHANGELOG.md                                     |  3 +++
 .../Private/Assert-AzureConnection.tests.ps1     |  3 ---
 .../Unit/Public/Search-AzResourceGraph.tests.ps1 | 10 ++++++++++
 4 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/AzResourceGraph/Public/Search-AzResourceGraph.ps1 b/AzResourceGraph/Public/Search-AzResourceGraph.ps1
index bb0c420..0bfd2fa 100644
--- a/AzResourceGraph/Public/Search-AzResourceGraph.ps1
+++ b/AzResourceGraph/Public/Search-AzResourceGraph.ps1
@@ -42,6 +42,9 @@ Allow partial scopes in the query. Only applicable for tenant and management gro
 Number of rows to request per page (1-1000).
 The function continues paging until all rows are retrieved.
 
+.PARAMETER Token
+Use to call Azure Resource Graph with a specified access token. Using this parameter will override any sign-in made with Connect-AzResourceGraph for a single command.
+
 .EXAMPLE
 # Execute a query stored in a file against two subscriptions
 Search-AzResourceGraph -QueryPath '.\vm-details.kql' `
@@ -90,7 +93,11 @@ function Search-AzResourceGraph {
         [Parameter(ParameterSetName = 'Path')]
         [Parameter(ParameterSetName = 'String')]
         [ValidateRange(1, 1000)]
-        [int]$PageSize = 1000
+        [int]$PageSize = 1000,
+
+        [Parameter(ParameterSetName = 'Path', DontShow)]
+        [Parameter(ParameterSetName = 'String', DontShow)]
+        [string]$Token
     )
 
     # Ensure only one of SubscriptionId or ManagementGroup is provided
@@ -98,7 +105,10 @@ function Search-AzResourceGraph {
         throw 'KQL Query can only be run against either a Subscription or a Management Group, not both.'
     }
 
-    Assert-AzureConnection -TokenSplat $script:TokenSplat
+    if (-not $PSBoundParameters.ContainsKey('Token')) {
+        Assert-AzureConnection -TokenSplat $script:TokenSplat
+        $Token = $script:Token.Token
+    }
 
     if ($PSCmdlet.ParameterSetName -eq 'Path') {
         $Query = Get-Content $QueryPath -Raw
@@ -120,7 +130,7 @@ function Search-AzResourceGraph {
     if ($PSBoundParameters.ContainsKey('ManagementGroup')) { $Body['managementGroups'] = @($ManagementGroup) }
 
     $Headers = @{
-        'Authorization' = "Bearer $($script:Token.Token)"
+        'Authorization' = "Bearer $Token"
         'Content-Type'  = 'application/json'
     }
 
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9dac62a..5ba7de3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Added Token parameter to Search-AzResourceGraph function to allow specifying a token for authentication
+
 ### Changed 
 - Updated dependency on AzAuth to version 2.5.0
 
diff --git a/tests/Unit/Private/Assert-AzureConnection.tests.ps1 b/tests/Unit/Private/Assert-AzureConnection.tests.ps1
index 9d2b18d..0d42c55 100644
--- a/tests/Unit/Private/Assert-AzureConnection.tests.ps1
+++ b/tests/Unit/Private/Assert-AzureConnection.tests.ps1
@@ -105,9 +105,6 @@ Describe 'Assert-AzureConnection' {
                     }
                 } | Should -Throw
             }
-
         }
     }
-
-
 }
\ No newline at end of file
diff --git a/tests/Unit/Public/Search-AzResourceGraph.tests.ps1 b/tests/Unit/Public/Search-AzResourceGraph.tests.ps1
index 08ffb55..d19f09c 100644
--- a/tests/Unit/Public/Search-AzResourceGraph.tests.ps1
+++ b/tests/Unit/Public/Search-AzResourceGraph.tests.ps1
@@ -54,4 +54,14 @@ Describe 'Search-AzResourceGraph' {
             $Body.managementGroups -contains 'myMG'
         }
     }
+
+    It 'Calls Assert-AzureConnection when not given a token' {
+        $null = Search-AzResourceGraph -QueryPath '~/foo.kql' -PageSize 10
+        Should -Invoke 'Assert-AzureConnection' -Times 1 -Exactly  -ModuleName 'AzResourceGraph'
+    }
+
+    It 'Does not call Assert-AzureConnection when given a token' {
+        $null = Search-AzResourceGraph -QueryPath '~/foo.kql' -PageSize 10 -Token 'fakeToken'
+        Should -Invoke 'Assert-AzureConnection' -Times 0 -Exactly  -ModuleName 'AzResourceGraph'
+    }
 }
\ No newline at end of file