Skip to content

Commit fff9504

Browse files
PrzemyslawKlysPrzemyslawKlys
committed
Address WAM review feedback
1 parent 79097a9 commit fff9504

6 files changed

Lines changed: 156 additions & 45 deletions

File tree

‎Private/Get-O365BrokerAccessToken.ps1‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,12 @@ function Get-O365BrokerAccessToken {
99
[Parameter(Mandatory, ParameterSetName = 'Scope')][string] $Scope,
1010
[string] $Tenant = 'organizations',
1111
[string] $ClientId = '1950a258-227b-4e31-a9cf-717495945fc2',
12+
[string] $Account,
1213
[switch] $ForcePrompt
1314
)
1415

15-
if ($PSEdition -ne 'Core') {
16-
throw 'MSAL WAM broker authentication requires PowerShell 7 or newer.'
16+
if ($PSEdition -ne 'Core' -or $PSVersionTable.PSVersion -lt [version] '7.4') {
17+
throw 'MSAL WAM broker authentication requires PowerShell 7.4 or newer.'
1718
}
1819

1920
$BrokerClientType = 'O365Essentials.Auth.BrokerTokenClient' -as [type]
@@ -23,9 +24,9 @@ function Get-O365BrokerAccessToken {
2324

2425
try {
2526
if ($PSCmdlet.ParameterSetName -eq 'Scope') {
26-
$Result = $BrokerClientType.GetMethod('AcquireTokenForScope').Invoke($null, @($Tenant, $Scope, $ClientId, [bool] $ForcePrompt))
27+
$Result = $BrokerClientType.GetMethod('AcquireTokenForScope').Invoke($null, @($Tenant, $Scope, $ClientId, $Account, [bool] $ForcePrompt))
2728
} else {
28-
$Result = $BrokerClientType.GetMethod('AcquireTokenForResource').Invoke($null, @($Tenant, $ResourceUrl, $ClientId, [bool] $ForcePrompt))
29+
$Result = $BrokerClientType.GetMethod('AcquireTokenForResource').Invoke($null, @($Tenant, $ResourceUrl, $ClientId, $Account, [bool] $ForcePrompt))
2930
}
3031
} catch [System.Reflection.TargetInvocationException] {
3132
$Inner = $_.Exception.InnerException

‎Public/Connect-O365Admin.ps1‎

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ function Connect-O365Admin {
8989
$PortalUserId = $null
9090
$PortalTenantId = $null
9191
$HeadersPortal = $null
92+
$RequestedUseWam = [bool] $UseWam
9293
if ($Headers) {
9394
if ($Headers.ExpiresOnUTC -gt [datetime]::UtcNow -and -not $ForceRefresh -and -not $HasPortalAttachInput) {
9495
Write-Verbose -Message "Connect-O365Admin - Using cache for connection $($Headers.UserName)"
@@ -102,7 +103,9 @@ function Connect-O365Admin {
102103
$Tenant = $Headers.Tenant
103104
$Subscription = $Headers.Subscription
104105
$RefreshToken = $Headers.RefreshToken
105-
$UseWam = $Headers.Contains('AuthenticationMode') -and $Headers.AuthenticationMode -eq 'WAM'
106+
if (-not $RequestedUseWam) {
107+
$UseWam = $Headers.Contains('AuthenticationMode') -and $Headers.AuthenticationMode -eq 'WAM'
108+
}
106109
if ($Headers.Contains('PortalWebSession')) { $PortalWebSession = $Headers.PortalWebSession }
107110
if ($Headers.Contains('AjaxSessionKey')) { $AjaxSessionKey = $Headers.AjaxSessionKey }
108111
if ($Headers.Contains('PortalRouteKey')) { $PortalRouteKey = $Headers.PortalRouteKey }
@@ -123,7 +126,9 @@ function Connect-O365Admin {
123126
$Tenant = $Script:AuthorizationO365Cache.Tenant
124127
$Subscription = $Script:AuthorizationO365Cache.Subscription
125128
$RefreshToken = $Script:AuthorizationO365Cache.RefreshToken
126-
$UseWam = $Script:AuthorizationO365Cache.Contains('AuthenticationMode') -and $Script:AuthorizationO365Cache.AuthenticationMode -eq 'WAM'
129+
if (-not $RequestedUseWam) {
130+
$UseWam = $Script:AuthorizationO365Cache.Contains('AuthenticationMode') -and $Script:AuthorizationO365Cache.AuthenticationMode -eq 'WAM'
131+
}
127132
if ($Script:AuthorizationO365Cache.Contains('PortalWebSession')) { $PortalWebSession = $Script:AuthorizationO365Cache.PortalWebSession }
128133
if ($Script:AuthorizationO365Cache.Contains('AjaxSessionKey')) { $AjaxSessionKey = $Script:AuthorizationO365Cache.AjaxSessionKey }
129134
if ($Script:AuthorizationO365Cache.Contains('PortalRouteKey')) { $PortalRouteKey = $Script:AuthorizationO365Cache.PortalRouteKey }
@@ -234,16 +239,17 @@ function Connect-O365Admin {
234239
# Read tenant information from the token so subsequent requests use the correct tenant
235240
$jwtTenant = if ($tokenGraph.id_token) { $tokenGraph.id_token } else { $tokenGraph.access_token }
236241
$tenantInfo = ConvertFrom-JSONWebToken -Token $jwtTenant
237-
if (-not $PSBoundParameters.ContainsKey('Tenant') -or $Tenant -eq 'organizations') {
238-
$Tenant = $tenantInfo.tid
239-
}
242+
if (-not $PSBoundParameters.ContainsKey('Tenant') -or $Tenant -eq 'organizations') {
243+
$Tenant = $tenantInfo.tid
244+
}
240245
$refresh = if ($UseWam) { $null } else { $tokenGraph.refresh_token }
246+
$WamAccount = if ($UseWam) { $tokenGraph.account } else { $null }
241247
$tokenO365 = $null
242248
$tokenAzure = $null
243249
try {
244250
Write-Verbose -Message "Connect-O365Admin - Acquiring token for admin.microsoft.com"
245251
if ($UseWam) {
246-
$tokenO365 = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl 'https://admin.microsoft.com/'
252+
$tokenO365 = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl 'https://admin.microsoft.com/' -Account $WamAccount
247253
} elseif ($PSCmdlet.ParameterSetName -eq 'App') {
248254
$tokenO365 = Get-O365OAuthToken -Tenant $Tenant -Scope $ScopesO365 -ClientId $ClientId -ClientSecret $ClientSecret -Certificate $Certificate -CertificatePassword $CertificatePassword
249255
} else {
@@ -253,7 +259,7 @@ function Connect-O365Admin {
253259
Write-Verbose -Message "Connect-O365Admin - admin.microsoft.com scope token failed. Falling back to portal resource audience. Error: $($_.Exception.Message)"
254260
try {
255261
if ($UseWam) {
256-
$tokenO365 = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl $ResourceAzure
262+
$tokenO365 = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl $ResourceAzure -Account $WamAccount
257263
} elseif ($PSCmdlet.ParameterSetName -eq 'App') {
258264
$tokenO365 = Get-O365OAuthToken -Tenant $Tenant -Resource $ResourceAzure -ClientId $ClientId -ClientSecret $ClientSecret -Certificate $Certificate -CertificatePassword $CertificatePassword
259265
} else {
@@ -269,7 +275,7 @@ function Connect-O365Admin {
269275
try {
270276
Write-Verbose -Message "Connect-O365Admin - Acquiring token for Teams (api.spaces.skype.com)"
271277
if ($UseWam) {
272-
$tokenTeams = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl 'https://api.spaces.skype.com/'
278+
$tokenTeams = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl 'https://api.spaces.skype.com/' -Account $WamAccount
273279
} elseif ($PSCmdlet.ParameterSetName -eq 'App') {
274280
$tokenTeams = Get-O365OAuthToken -Tenant $Tenant -Scope $ScopesTeams -ClientId $ClientId -ClientSecret $ClientSecret -Certificate $Certificate -CertificatePassword $CertificatePassword
275281
} else {
@@ -283,7 +289,7 @@ function Connect-O365Admin {
283289
try {
284290
Write-Verbose -Message "Connect-O365Admin - Acquiring token for Azure"
285291
if ($UseWam) {
286-
$tokenAzure = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl $ResourceAzure
292+
$tokenAzure = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl $ResourceAzure -Account $WamAccount
287293
} elseif ($PSCmdlet.ParameterSetName -eq 'App') {
288294
$tokenAzure = Get-O365OAuthToken -Tenant $Tenant -Resource $ResourceAzure -ClientId $ClientId -ClientSecret $ClientSecret -Certificate $Certificate -CertificatePassword $CertificatePassword
289295
} else {
@@ -297,7 +303,7 @@ function Connect-O365Admin {
297303
try {
298304
Write-Verbose -Message "Connect-O365Admin - Acquiring token for Azure management"
299305
if ($UseWam) {
300-
$tokenARM = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl 'https://management.azure.com/'
306+
$tokenARM = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl 'https://management.azure.com/' -Account $WamAccount
301307
} elseif ($PSCmdlet.ParameterSetName -eq 'App') {
302308
$tokenARM = Get-O365OAuthToken -Tenant $Tenant -Scope $ScopesARM -ClientId $ClientId -ClientSecret $ClientSecret -Certificate $Certificate -CertificatePassword $CertificatePassword
303309
} else {
@@ -333,9 +339,9 @@ function Connect-O365Admin {
333339
Write-Verbose -Message "Connect-O365Admin - Acquiring token for Substrate using $($attempt.type): $($attempt.value)"
334340
if ($UseWam) {
335341
if ($attempt.type -eq 'scope') {
336-
$tokenSubstrate = Get-O365BrokerAccessToken -Tenant $Tenant -Scope $attempt.value
342+
$tokenSubstrate = Get-O365BrokerAccessToken -Tenant $Tenant -Scope $attempt.value -Account $WamAccount
337343
} else {
338-
$tokenSubstrate = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl $attempt.value
344+
$tokenSubstrate = Get-O365BrokerAccessToken -Tenant $Tenant -ResourceUrl $attempt.value -Account $WamAccount
339345
}
340346
} elseif ($attempt.type -eq 'resource') {
341347
if ($PSCmdlet.ParameterSetName -eq 'App') {

‎README.MD‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -76,10 +76,10 @@ If we have MFA we simply let it query you for MFA code.
7676
$null = Connect-O365Admin -Verbose
7777
```
7878

79-
On Windows PowerShell 7+, you can opt into the Web Account Manager broker through
80-
the bundled MSAL helper. This allows Windows Hello, Conditional Access, FIDO, and
81-
the Windows account picker to participate in delegated sign-in while the MSAL
82-
assemblies stay isolated inside the packaged module:
79+
On Windows PowerShell 7.4+, you can opt into the Web Account Manager broker
80+
through the bundled MSAL helper. This allows Windows Hello, Conditional Access,
81+
FIDO, and the Windows account picker to participate in delegated sign-in while
82+
the MSAL assemblies stay isolated inside the packaged module:
8383

8484
```powershell
8585
$null = Connect-O365Admin -UseWam -Tenant 'contoso.onmicrosoft.com' -Verbose

‎Sources/O365Essentials.Auth/BrokerTokenClient.cs‎

Lines changed: 78 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,18 @@ public static BrokerTokenResult AcquireTokenForResource(
2525
string? tenant,
2626
string resourceUrl,
2727
string clientId,
28+
string? accountUsername,
2829
bool forcePrompt)
2930
{
3031
var scopeCandidates = BuildScopeCandidatesFromResource(resourceUrl);
31-
return AcquireTokenAsync(tenant, scopeCandidates, clientId, forcePrompt).GetAwaiter().GetResult();
32+
return AcquireTokenAsync(tenant, scopeCandidates, clientId, accountUsername, forcePrompt).GetAwaiter().GetResult();
3233
}
3334

3435
public static BrokerTokenResult AcquireTokenForScope(
3536
string? tenant,
3637
string scope,
3738
string clientId,
39+
string? accountUsername,
3840
bool forcePrompt)
3941
{
4042
var scopes = SplitScopes(scope);
@@ -43,13 +45,14 @@ public static BrokerTokenResult AcquireTokenForScope(
4345
throw new ArgumentException("At least one scope must be provided.", nameof(scope));
4446
}
4547

46-
return AcquireTokenAsync(tenant, new[] { scopes }, clientId, forcePrompt).GetAwaiter().GetResult();
48+
return AcquireTokenAsync(tenant, new[] { scopes }, clientId, accountUsername, forcePrompt).GetAwaiter().GetResult();
4749
}
4850

4951
private static async Task<BrokerTokenResult> AcquireTokenAsync(
5052
string? tenant,
5153
IReadOnlyList<string[]> scopeCandidates,
5254
string clientId,
55+
string? accountUsername,
5356
bool forcePrompt)
5457
{
5558
if (!OperatingSystem.IsWindows())
@@ -86,7 +89,8 @@ private static async Task<BrokerTokenResult> AcquireTokenAsync(
8689
{
8790
try
8891
{
89-
var result = await AcquireTokenForScopesAsync(application, scopes, forcePrompt).ConfigureAwait(false);
92+
var result = await AcquireTokenForScopesAsync(application, scopes, accountUsername, forcePrompt).ConfigureAwait(false);
93+
EnsureAccountMatches(result, accountUsername);
9094
return ToResult(result);
9195
}
9296
catch (MsalException ex) when (CanTryNextScopeCandidate(ex))
@@ -105,12 +109,20 @@ private static async Task<BrokerTokenResult> AcquireTokenAsync(
105109
private static async Task<AuthenticationResult> AcquireTokenForScopesAsync(
106110
IPublicClientApplication application,
107111
string[] scopes,
112+
string? accountUsername,
108113
bool forcePrompt)
109114
{
115+
var expectedUsername = string.IsNullOrWhiteSpace(accountUsername) ? null : accountUsername.Trim();
110116
if (!forcePrompt)
111117
{
112118
var accounts = await application.GetAccountsAsync().ConfigureAwait(false);
113-
foreach (var account in accounts)
119+
var candidateAccounts = string.IsNullOrWhiteSpace(expectedUsername)
120+
? accounts
121+
: accounts.Where(account =>
122+
!string.IsNullOrWhiteSpace(account.Username)
123+
&& account.Username.Equals(expectedUsername, StringComparison.OrdinalIgnoreCase));
124+
125+
foreach (var account in candidateAccounts)
114126
{
115127
try
116128
{
@@ -121,20 +133,29 @@ private static async Task<AuthenticationResult> AcquireTokenForScopesAsync(
121133
}
122134
}
123135

124-
try
125-
{
126-
return await application
127-
.AcquireTokenSilent(scopes, PublicClientApplication.OperatingSystemAccount)
128-
.ExecuteAsync()
129-
.ConfigureAwait(false);
130-
}
131-
catch (MsalUiRequiredException)
136+
if (string.IsNullOrWhiteSpace(expectedUsername))
132137
{
138+
try
139+
{
140+
return await application
141+
.AcquireTokenSilent(scopes, PublicClientApplication.OperatingSystemAccount)
142+
.ExecuteAsync()
143+
.ConfigureAwait(false);
144+
}
145+
catch (MsalUiRequiredException)
146+
{
147+
}
133148
}
134149
}
135150

136-
return await application
137-
.AcquireTokenInteractive(scopes)
151+
var interactiveBuilder = application
152+
.AcquireTokenInteractive(scopes);
153+
if (!string.IsNullOrWhiteSpace(expectedUsername))
154+
{
155+
interactiveBuilder = interactiveBuilder.WithLoginHint(expectedUsername);
156+
}
157+
158+
return await interactiveBuilder
138159
.WithPrompt(Prompt.SelectAccount)
139160
.ExecuteAsync()
140161
.ConfigureAwait(false);
@@ -193,12 +214,43 @@ private static bool CanTryNextScopeCandidate(MsalException exception)
193214
|| exception.ErrorCode.Equals("invalid_request", StringComparison.OrdinalIgnoreCase);
194215
}
195216

217+
private static void EnsureAccountMatches(AuthenticationResult result, string? expectedUsername)
218+
{
219+
if (string.IsNullOrWhiteSpace(expectedUsername) || result.Account is null)
220+
{
221+
return;
222+
}
223+
224+
if (string.IsNullOrWhiteSpace(result.Account.Username)
225+
|| !result.Account.Username.Equals(expectedUsername.Trim(), StringComparison.OrdinalIgnoreCase))
226+
{
227+
throw new InvalidOperationException(
228+
$"MSAL WAM returned account '{result.Account.Username}' while O365Essentials expected '{expectedUsername}'.");
229+
}
230+
}
231+
196232
private static IntPtr GetConsoleOrTerminalWindow()
197233
{
198234
var consoleHandle = GetConsoleWindow();
199-
return consoleHandle == IntPtr.Zero
200-
? IntPtr.Zero
201-
: GetAncestor(consoleHandle, GetAncestorFlags.GetRootOwner);
235+
if (consoleHandle != IntPtr.Zero)
236+
{
237+
var owner = GetAncestor(consoleHandle, GetAncestorFlags.GetRootOwner);
238+
return owner == IntPtr.Zero ? consoleHandle : owner;
239+
}
240+
241+
var foregroundHandle = GetForegroundWindow();
242+
if (foregroundHandle != IntPtr.Zero)
243+
{
244+
return foregroundHandle;
245+
}
246+
247+
var shellHandle = GetShellWindow();
248+
if (shellHandle != IntPtr.Zero)
249+
{
250+
return shellHandle;
251+
}
252+
253+
return GetDesktopWindow();
202254
}
203255

204256
private enum GetAncestorFlags
@@ -209,6 +261,15 @@ private enum GetAncestorFlags
209261
[DllImport("kernel32.dll")]
210262
private static extern IntPtr GetConsoleWindow();
211263

264+
[DllImport("user32.dll")]
265+
private static extern IntPtr GetForegroundWindow();
266+
267+
[DllImport("user32.dll")]
268+
private static extern IntPtr GetShellWindow();
269+
270+
[DllImport("user32.dll")]
271+
private static extern IntPtr GetDesktopWindow();
272+
212273
[DllImport("user32.dll", ExactSpelling = true)]
213274
private static extern IntPtr GetAncestor(IntPtr hwnd, GetAncestorFlags flags);
214275
}

‎Tests/Get-O365OAuthToken.Tests.ps1‎

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ Describe 'Connect-O365Admin portal token' {
152152
}
153153
}
154154
Mock -ModuleName O365Essentials Get-O365BrokerAccessToken -MockWith {
155-
param($Tenant, $ResourceUrl, $Scope, $ForcePrompt)
155+
param($Tenant, $ResourceUrl, $Scope, $Account, $ForcePrompt)
156156
$Target = if ($ResourceUrl) { $ResourceUrl } else { $Scope }
157157
[pscustomobject]@{
158158
access_token = "token:$Target"
@@ -172,7 +172,51 @@ Describe 'Connect-O365Admin portal token' {
172172
$result.HeadersO365.Authorization | Should -Be 'Bearer token:https://admin.microsoft.com/'
173173
Assert-MockCalled Get-O365OAuthToken -ModuleName O365Essentials -Exactly 0
174174
Assert-MockCalled Get-O365BrokerAccessToken -ModuleName O365Essentials -ParameterFilter {
175-
$ResourceUrl -eq 'https://admin.microsoft.com/'
175+
$ResourceUrl -eq 'https://admin.microsoft.com/' -and $Account -eq 'user@contoso.com'
176+
} -Exactly 1
177+
}
178+
179+
It 'honors explicit WAM refresh when an expired OAuth cache exists' {
180+
InModuleScope O365Essentials {
181+
$script:AuthorizationO365Cache = [ordered] @{
182+
Credential = $null
183+
ClientId = $null
184+
ClientSecret = $null
185+
Certificate = $null
186+
CertificatePassword = $null
187+
AuthenticationMode = 'OAuth'
188+
UserName = 'old@contoso.com'
189+
Tenant = 'tenant-id'
190+
Subscription = $null
191+
RefreshToken = 'old-refresh-token'
192+
ExpiresOnUTC = ([datetime]::UtcNow).AddMinutes(-5)
193+
}
194+
}
195+
Mock -ModuleName O365Essentials ConvertFrom-JSONWebToken -MockWith {
196+
[pscustomobject]@{
197+
tid = 'tenant-id'
198+
upn = 'user@contoso.com'
199+
}
200+
}
201+
Mock -ModuleName O365Essentials Get-O365BrokerAccessToken -MockWith {
202+
param($Tenant, $ResourceUrl, $Scope, $Account, $ForcePrompt)
203+
$Target = if ($ResourceUrl) { $ResourceUrl } else { $Scope }
204+
[pscustomobject]@{
205+
access_token = "token:$Target"
206+
expires_on = ([datetime]::UtcNow).AddHours(1)
207+
tenant_id = 'tenant-id'
208+
account = 'user@contoso.com'
209+
}
210+
}
211+
Mock -ModuleName O365Essentials Get-O365OAuthToken -MockWith { throw 'legacy OAuth path should not be used' }
212+
213+
$result = Connect-O365Admin -UseWam -ForceRefresh
214+
215+
$result.AuthenticationMode | Should -Be 'WAM'
216+
$result.UserName | Should -Be 'user@contoso.com'
217+
Assert-MockCalled Get-O365OAuthToken -ModuleName O365Essentials -Exactly 0
218+
Assert-MockCalled Get-O365BrokerAccessToken -ModuleName O365Essentials -ParameterFilter {
219+
$ResourceUrl -eq 'https://graph.microsoft.com/' -and $ForcePrompt
176220
} -Exactly 1
177221
}
178222

0 commit comments

Comments
 (0)