From 85255a7ddef6d6e57d9f9b87bd081c3b7a0fbb9b Mon Sep 17 00:00:00 2001
From: Jiri Formacek <jiri.formacek@greycorbel.com>
Date: Thu, 17 Oct 2024 08:53:03 +0200
Subject: [PATCH] fix: tenant ID handling in New-ExoConnection

---
 Commands/Public/New-ExoConnection.ps1 | 17 ++++++++++++++---
 Module/ExoHelper/ExoHelper.psd1       |  2 +-
 Module/ExoHelper/ExoHelper.psm1       | 17 ++++++++++++++---
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/Commands/Public/New-ExoConnection.ps1 b/Commands/Public/New-ExoConnection.ps1
index 5d9f448..d5a3679 100644
--- a/Commands/Public/New-ExoConnection.ps1
+++ b/Commands/Public/New-ExoConnection.ps1
@@ -65,8 +65,19 @@ param
             IsIPPS = $IPPS.IsPresent
             HttpClient = new-object System.Net.Http.HttpClient
         }
-        $claims = Get-ExoToken -Connection $Connection | Test-AadToken -PayloadOnly
-        $Connection.TenantId = $claims.tid
+        
+        #explicitly authenticate when establishing connection to catch any authentication problems early
+        Get-ExoToken -Connection $Connection | Out-Null
+        if([string]::IsNullOrEmpty($TenantId))
+        {
+            $TenantId = $AuthenticationFactory.TenantId
+        }
+        if([string]::IsNullOrEmpty($TenantId))
+        {
+            throw (new-object ExoHelper.ExoException([System.Net.HttpStatusCode]::BadRequest, 'ExoMissingTenantId', 'ExoInitializationError', 'TenantId is not specified and cannot be determined automatically - please specify TenantId parameter'))
+        }
+        $Connection.TenantId = $TenantId
+        
         if($IPPS)
         {
             $Connection.ConnectionUri = "https://eur02b.ps.compliance.protection.outlook.com/adminapi/beta/$($Connection.TenantId)/InvokeCommand"
@@ -86,7 +97,7 @@ param
             else
             {
                 #likely app-only context - use same static anchor mailbox as ExchangeOnlineManagement module uses
-                $Connection.AnchorMailbox = "DiscoverySearchMailbox{D919BA05-46A6-415f-80AD-7E09334BB852}@$tenantId"
+                $Connection.AnchorMailbox = "UPN:DiscoverySearchMailbox{D919BA05-46A6-415f-80AD-7E09334BB852}@$tenantId"
             }
         }
         else
diff --git a/Module/ExoHelper/ExoHelper.psd1 b/Module/ExoHelper/ExoHelper.psd1
index d49a67a..315faf2 100644
--- a/Module/ExoHelper/ExoHelper.psd1
+++ b/Module/ExoHelper/ExoHelper.psd1
@@ -12,7 +12,7 @@
 RootModule = 'ExoHelper.psm1'
 
 # Version number of this module.
-ModuleVersion = '3.0.0'
+ModuleVersion = '3.0.1'
 
 # Supported PSEditions
 CompatiblePSEditions = @('Desktop', 'Core')
diff --git a/Module/ExoHelper/ExoHelper.psm1 b/Module/ExoHelper/ExoHelper.psm1
index 049ea26..001646e 100644
--- a/Module/ExoHelper/ExoHelper.psm1
+++ b/Module/ExoHelper/ExoHelper.psm1
@@ -364,8 +364,19 @@ param
             IsIPPS = $IPPS.IsPresent
             HttpClient = new-object System.Net.Http.HttpClient
         }
-        $claims = Get-ExoToken -Connection $Connection | Test-AadToken -PayloadOnly
-        $Connection.TenantId = $claims.tid
+        
+        #explicitly authenticate when establishing connection to catch any authentication problems early
+        Get-ExoToken -Connection $Connection | Out-Null
+        if([string]::IsNullOrEmpty($TenantId))
+        {
+            $TenantId = $AuthenticationFactory.TenantId
+        }
+        if([string]::IsNullOrEmpty($TenantId))
+        {
+            throw (new-object ExoHelper.ExoException([System.Net.HttpStatusCode]::BadRequest, 'ExoMissingTenantId', 'ExoInitializationError', 'TenantId is not specified and cannot be determined automatically - please specify TenantId parameter'))
+        }
+        $Connection.TenantId = $TenantId
+        
         if($IPPS)
         {
             $Connection.ConnectionUri = "https://eur02b.ps.compliance.protection.outlook.com/adminapi/beta/$($Connection.TenantId)/InvokeCommand"
@@ -385,7 +396,7 @@ param
             else
             {
                 #likely app-only context - use same static anchor mailbox as ExchangeOnlineManagement module uses
-                $Connection.AnchorMailbox = "DiscoverySearchMailbox{D919BA05-46A6-415f-80AD-7E09334BB852}@$tenantId"
+                $Connection.AnchorMailbox = "UPN:DiscoverySearchMailbox{D919BA05-46A6-415f-80AD-7E09334BB852}@$tenantId"
             }
         }
         else