From 41d2bdfc15e5f98462dc8f8538acb269f7acfc7e Mon Sep 17 00:00:00 2001
From: BradleyBartlett <bradley.bartlett@microsoft.com>
Date: Thu, 21 Sep 2017 18:19:31 -0700
Subject: [PATCH] fixes to custom RBAC role workaround; token collection
 cleaned up

---
 Registration/RegisterWithAzure.psm1 | 101 ++++++++++------------------
 1 file changed, 35 insertions(+), 66 deletions(-)

diff --git a/Registration/RegisterWithAzure.psm1 b/Registration/RegisterWithAzure.psm1
index 5a265a38..acd38782 100644
--- a/Registration/RegisterWithAzure.psm1
+++ b/Registration/RegisterWithAzure.psm1
@@ -736,24 +736,28 @@ function New-RBACAssignment{
     if (-not $customRoleDefined)
     {
         $customRoleName = "Registration Reader-$($RegistrationResource.SubscriptionId)"
-        # Create new RBAC role definition
-        $role = Get-AzureRmRoleDefinition -Name 'Reader'
-        $role.Name = $customRoleName
-        $role.id = [guid]::newguid()
-        $role.IsCustom = $true
-        $role.Actions.Add('Microsoft.AzureStack/registrations/products/listDetails/action')
-        $role.Actions.Add('Microsoft.AzureStack/registrations/products/read')
-        $role.AssignableScopes.Clear()
-        $role.AssignableScopes.Add("/subscriptions/$($RegistrationResource.SubscriptionId)")
-        $role.Description = "Custom RBAC role for registration actions such as downloading products from Azure marketplace"
-        try
-        {
-            New-AzureRmRoleDefinition -Role $role
-        }
-        catch
+        $customRoleDefined = Get-AzureRmRoleDefinition -Name $customRoleName
+        if (-not $customRoleDefined)
         {
-            Log-Throw -Message "Defining custom RBAC role $customRoleName failed: `r`n$($_.Exception)" -CallingFunction $PSCmdlet.MyInvocation.InvocationName
-        }
+            # Create new RBAC role definition
+            $role = Get-AzureRmRoleDefinition -Name 'Reader'
+            $role.Name = $customRoleName
+            $role.id = [guid]::newguid()
+            $role.IsCustom = $true
+            $role.Actions.Add('Microsoft.AzureStack/registrations/products/listDetails/action')
+            $role.Actions.Add('Microsoft.AzureStack/registrations/products/read')
+            $role.AssignableScopes.Clear()
+            $role.AssignableScopes.Add("/subscriptions/$($RegistrationResource.SubscriptionId)")
+            $role.Description = "Custom RBAC role for registration actions such as downloading products from Azure marketplace"
+            try
+            {
+                New-AzureRmRoleDefinition -Role $role
+            }
+            catch
+            {
+                Log-Throw -Message "Defining custom RBAC role $customRoleName failed: `r`n$($_.Exception)" -CallingFunction $PSCmdlet.MyInvocation.InvocationName
+            }
+        }        
     }
 
     # Determine if custom RBAC role has been assigned
@@ -827,7 +831,8 @@ function Connect-AzureAccount{
             Add-AzureRmAccount -SubscriptionId $SubscriptionId       
             Set-AzureRmContext -SubscriptionId $SubscriptionId -TenantId $AzureDirectoryTenantId
             $environment = Get-AzureRmEnvironment -Name $AzureEnvironmentName
-            $subscription = Get-AzureRmSubscription -SubscriptionId $SubscriptionId 
+            $subscription = Get-AzureRmSubscription -SubscriptionId $SubscriptionId
+            $context = Get-AzureRmContext
         }
         catch
         {
@@ -840,58 +845,22 @@ function Connect-AzureAccount{
     }
 
 
-    [Version]$azurePSVersion = (Get-Module AzureRm.Profile).Version
-    Log-Output "Using AzureRm.Profile version: $azurePSVersion"
-
-    if ($azurePSVersion -ge [Version]"3.3.2")
-    {
-        $tokens = [Microsoft.Azure.Commands.Common.Authentication.AzureSession]::Instance.TokenCache.ReadItems()
-        if (-not $tokens -or ($tokens.Count -le 0))
-        {
-            $tokens = $context.TokenCache.ReadItems()
+    $tokens = @()
+    try{$tokens += [Microsoft.Azure.Commands.Common.Authentication.AzureSession]::Instance.TokenCache.ReadItems()}catch{}
+    try{$tokens += [Microsoft.IdentityModel.Clients.ActiveDirectory.TokenCache]::DefaultShared.ReadItems()}catch{}
+    try{$tokens += $context.TokenCache.ReadItems()}catch{}
 
-            if (-not $tokens -or ($tokens.Count -le 0))
-            {
-                Log-Throw -Message "Token cache is empty `r`n$($_.Exception)" -CallingFunction $PSCmdlet.MyInvocation.InvocationName
-            }
-            else
-            {
-                $token = $tokens |
-                    Where Resource -EQ $environment.ActiveDirectoryServiceEndpointResourceId |
-                    Where { $_.TenantId -eq $subscription.TenantId } |
-                    Where { $_.ExpiresOn -gt [datetime]::UtcNow } |
-                    Select -First 1
-            }
-        }
-        else
-        {
-            $token = $tokens |
-                Where Resource -EQ $environment.ActiveDirectoryServiceEndpointResourceId |
-                Where { $_.TenantId -eq $subscription.TenantId } |
-                Where { $_.ExpiresOn -gt [datetime]::UtcNow } |
-                Select -First 1
-        }
-    }
-    else
+    if (-not $tokens -or ($tokens.Count -le 0))
     {
-        $tokens = [Microsoft.IdentityModel.Clients.ActiveDirectory.TokenCache]::DefaultShared.ReadItems()
-        if (-not $tokens -or ($tokens.Count -le 0))
-        {            
-            if (-not $tokens -or ($tokens.Count -le 0))
-            {
-                Log-Throw -Message "Token cache is empty `r`n$($_.Exception)" -CallingFunction $PSCmdlet.MyInvocation.InvocationName
-            }
-        }
-        else
-        {
-            $token = $tokens |
-                Where Resource -EQ $environment.ActiveDirectoryServiceEndpointResourceId |
-                Where { $_.TenantId -eq $subscription.TenantId } |
-                Where { $_.ExpiresOn -gt [datetime]::UtcNow } |
-                Select -First 1
-        }
+        Log-Throw -Message "Token cache is empty `r`n$($_.Exception)" -CallingFunction $PSCmdlet.MyInvocation.InvocationName
     }
 
+    $token = $tokens |
+        Where Resource -EQ $environment.ActiveDirectoryServiceEndpointResourceId |
+        Where { $_.TenantId -eq $subscription.TenantId } |
+        Sort ExpiresOn |
+        Select -Last 1
+
 
     if (-not $token)
     {