functions/Switch-AzContext.ps1

function Switch-AzContext {
    # The following SuppressMessageAttribute entries are used to surpress
    # PSScriptAnalyzer tests against known exceptions as per:
    # https://github.com/powershell/psscriptanalyzer#suppressing-rules
    # CredentialsDataFile is not a password
    [Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSAvoidUsingPlainTextForPassword', 'CredentialsDataFile')]
    # CredentialsJsonFile is not a password
    [Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSAvoidUsingPlainTextForPassword', 'CredentialsJsonFile')]
    # CredentialsObject is not a password
    [Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSAvoidUsingPlainTextForPassword', 'CredentialsObject')]
    # Looking for secure alternative, but the credentials are stored in plain text
    [Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSAvoidUsingConvertToSecureStringWithPlainText', '')]
    [CmdletBinding()]
    param (

        [Parameter(Mandatory = $true, ParameterSetName = 'ListAvailableFromDataFile')]
        [Parameter(Mandatory = $true, ParameterSetName = 'ListAvailableFromObject')]
        [switch]$ListAvailable,

        [Parameter(Mandatory = $true, ParameterSetName = 'FriendlyNameFromDataFile')]
        [Parameter(Mandatory = $true, ParameterSetName = 'FriendlyNameFromObject')]
        [Alias("FriendlyName")]
        [string]$Name,

        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromDataFile')]
        [Parameter(Mandatory = $true, ParameterSetName = 'TenantIdFromDataFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromObject')]
        [Parameter(Mandatory = $true, ParameterSetName = 'TenantIdFromObject')]
        [Parameter(Mandatory = $false, ParameterSetName = 'UserContext')]
        [string]$TenantId,

        [Parameter(Mandatory = $false, ParameterSetName = 'UserContext')]
        [switch]$UserContext,

        [Parameter(Mandatory = $true, ParameterSetName = 'FriendlyNameFromDataFile')]
        [Parameter(Mandatory = $true, ParameterSetName = 'TenantIdFromDataFile')]
        [Parameter(Mandatory = $true, ParameterSetName = 'ListAvailableFromDataFile')]
        [ValidateScript( { $_ | Test-Path -IsValid -PathType Container })]
        [string]$CredentialsDataFile,

        [Parameter(Mandatory = $true, ParameterSetName = 'FriendlyNameFromJsonFile')]
        [Parameter(Mandatory = $true, ParameterSetName = 'TenantIdFromJsonFile')]
        [Parameter(Mandatory = $true, ParameterSetName = 'ListAvailableFromJsonFile')]
        [ValidateScript( { $_ | Test-Path -IsValid -PathType Container })]
        [string]$CredentialsJsonFile,

        [Parameter(Mandatory = $true, ValueFromPipeline, ParameterSetName = 'FriendlyNameFromObject')]
        [Parameter(Mandatory = $true, ValueFromPipeline, ParameterSetName = 'TenantIdFromObject')]
        [Parameter(Mandatory = $true, ValueFromPipeline, ParameterSetName = 'ListAvailableFromObject')]
        [object]$CredentialsObject,

        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromDataFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromDataFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromJsonFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromJsonFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromObject')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromObject')]
        [switch]$UseAzLogin,

        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromDataFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromDataFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromJsonFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromJsonFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromObject')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromObject')]
        [switch]$UseTerraform,

        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromDataFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromDataFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromJsonFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromJsonFile')]
        [Parameter(Mandatory = $false, ParameterSetName = 'FriendlyNameFromObject')]
        [Parameter(Mandatory = $false, ParameterSetName = 'TenantIdFromObject')]
        [switch]$UseDefaultSubscription
    )

    begin {

        if ($UseDefaultSubscription) {
            Write-Warning "User has selected [UseDefaultSubscription]. Feature pending development."
        }

        function Get-CredentialsList {
            [CmdletBinding()]
            [OutputType([array])]
            param (
                [Parameter(Mandatory = $true, ValueFromPipeline, Position = 0)]
                [object]$InputObject
            )
            process {
                # Initialize empty array to store available credentials
                $AvailableCredentialBlocks = @()
                # Check
                if (-not (Test-IsCredentialsBlock $InputObject)) {
                    foreach ($Key in $InputObject.Keys) {
                        if (Test-IsCredentialsBlock $InputObject."$Key") {
                            $AvailableCredentialBlocks += [PSCustomObject]@{
                                name     = $Key
                                clientId = $InputObject."$Key".clientId
                                tenantId = $InputObject."$Key".tenantId
                            }
                        }
                        else {
                            Write-Warning "The provided InputObject contains an item [$Key] which isn't a valid CredentialBlock."
                        }
                    }
                }
                else {
                    $AvailableCredentialBlocks += [PSCustomObject]@{
                        name     = "$($InputObject.name ? $InputObject.name : "<none>")"
                        clientId = "$InputObject.clientId"
                        tenantId = "$InputObject.tenantId"
                    }
                }
                # Return available credentials if found
                if ($AvailableCredentialBlocks.Count -ge 1) {
                    return $AvailableCredentialBlocks
                }
                else {
                    return $null
                }
            }

        }

        function Test-IsCredentialsBlock {
            [CmdletBinding()]
            [OutputType([bool], [array])]
            param (
                [Parameter(Mandatory = $true, ValueFromPipeline, Position = 0)]
                [object]$InputObject,
                [Parameter(Mandatory = $false)]
                [switch]$ListMissingProperties
            )
            begin {
                $MandatoryProperties = @(
                    "tenantId"
                    "clientId"
                    "clientSecret"
                )
            }
            process {
                $MissingProperties = @()
                foreach ($Property in $MandatoryProperties) {
                    if ($Property -notin $InputObject.Keys) {
                        $MissingProperties += $Property
                    }
                }
                $MissingPropertiesCount = $MissingProperties.Count
                if ($ListMissingProperties) {
                    return $MissingProperties
                }
                elseif ($MissingPropertiesCount -eq 0) {
                    return $true
                }
                else {
                    return $false
                }
            }
        }

        function Get-CredentialsBlockByName {
            [CmdletBinding()]
            [OutputType([object])]
            param (
                [Parameter(Mandatory = $true, ValueFromPipeline, Position = 0)]
                [object]$InputObject,
                [Parameter(Mandatory = $true, Position = 1)]
                [string]$Name,
                [Parameter(Mandatory = $false, Position = 2)]
                [string]$TenantId
            )
            process {
                if (-not (Test-IsCredentialsBlock $InputObject)) {
                    # Check whether InputObject contains CredentialBlock matching Name
                    $InputObjectKeySearch = $InputObject.Keys -match $Name
                    $CountKeySearchMatches = $InputObjectKeySearch.Count
                    # Return match if found
                    if ($CountKeySearchMatches -eq 1) {
                        if ($TenantId) {
                            try {
                                Get-CredentialsBlockByTenantId $InputObject."$InputObjectKeySearch" -TenantId $TenantId -ErrorAction Stop
                            }
                            catch {
                                throw "Credentials for Name [$Name] do not match TenantId [$TenantId] in InputObject. Please check the provided TenantId."
                            }
                        }
                        else {
                            return $InputObject."$InputObjectKeySearch"
                        }
                    }
                    # Throw error if multiple matches found and not resolved by TenantId
                    elseif ($CountKeySearchMatches -gt 1) {
                        # Try to filter results using TenantId if provided
                        if ($TenantId) {
                            try {
                                Get-CredentialsBlockByTenantId $InputObject -TenantId $TenantId -ErrorAction Stop
                            }
                            catch {
                                throw "Found [$CountKeySearchMatches] matches for Name [$Name] and TenantId [$TenantId] in InputObject. Please provide a more specific name, or ensure no duplicates exist. The following results were returned:`n $($InputObjectKeySearch -join "`n ")"
                            }
                        }
                        else {
                            throw "Found [$CountKeySearchMatches] matches for Name [$Name] in InputObject. Please provide a more specific name, or ensure no duplicates exist. The following results were returned:`n $($InputObjectKeySearch -join "`n ")"
                        }
                    }
                    # Throw error if no matches found
                    elseif ($CountKeySearchMatches -lt 1) {
                        throw "Found [0] matches for Name [$Name] in InputObject."
                    }
                    # Throw unexpected error (should never occur)
                    else {
                        throw "Unexpected error occurred."
                    }
                }
                else {
                    if ($TenantId) {
                        try {
                            Get-CredentialsBlockByTenantId $InputObject -TenantId $TenantId -ErrorAction Stop
                        }
                        catch {
                            throw "CredentialsBlock for Name [$Name] does not match TenantId [$TenantId] in InputObject. Please check the provided TenantId."
                        }
                    }
                    else {
                        return $InputObject
                    }
                }
            }
        }

        function Get-CredentialsBlockByTenantId {
            [CmdletBinding()]
            [OutputType([object])]
            param (
                [Parameter(Mandatory = $true, ValueFromPipeline, Position = 0)]
                [object]$InputObject,
                [Parameter(Mandatory = $true, Position = 1)]
                [string]$TenantId
            )
            process {
                # Check whether InputObject contains CredentialBlock containing TenantId
                $InputObjectTenantSearch = @()
                if (-not (Test-IsCredentialsBlock $InputObject)) {
                    foreach ($Key in $InputObject.Keys) {
                        $CredentialsBlock = $InputObject."$Key"
                        $IsCredentialsBlock = Test-IsCredentialsBlock $CredentialsBlock
                        if ($IsCredentialsBlock) {
                            $IsTenantIdMatch = ($CredentialsBlock.TenantId -ccontains $TenantId)
                        }
                        if ($IsTenantIdMatch) {
                            $InputObjectTenantSearch += $CredentialsBlock
                        }
                    }
                }
                else {
                    $IsTenantIdMatch = ($InputObject.TenantId -ccontains $TenantId)
                    if ($IsTenantIdMatch) {
                        $InputObjectTenantSearch += $InputObject
                    }
                }
                $CountTenantSearchMatches = $InputObjectTenantSearch.Count
                # Return match if found
                if ($CountTenantSearchMatches -eq 1) {
                    return $InputObjectTenantSearch
                }
                # Throw error if multiple matches found
                elseif ($CountTenantSearchMatches -gt 1) {
                    throw "Found [$CountTenantSearchMatches] matches for TenantId [$TenantId] in InputObject. Please filter by name, or ensure no duplicates exist."
                }
                # Throw error if no matches found
                elseif ($CountKeySearchMatches -lt 1) {
                    throw "Found [$CountKeySearchMatches] matches for TenantId [$TenantId] in InputObject."
                }
                # Throw unexpected error (should never occur)
                else {
                    throw "Unexpected error occurred."
                }
            }
        }

    }

    process {

        ##############################################################
        # Process logon request for ['UserContext'] with [-TenantId] #
        ##############################################################

        if ($UserContext -and $TenantId) {
            Write-Verbose "Setting context using user credentials and TenantId [$TenantId]"
            # First, clear existing Azure context
            Clear-AzContext -Force | Out-Null
            $ctx = Connect-AzAccount -UseDeviceAuthentication -Tenant $TenantId
        }

        #################################################################
        # Process logon request for ['UserContext'] without [-TenantId] #
        #################################################################

        elseif ($UserContext) {
            Write-Verbose "Setting context using user credentials"
            # First, clear existing Azure context
            Clear-AzContext -Force | Out-Null
            $ctx = Connect-AzAccount -UseDeviceAuthentication
        }

        ###################################################################################
        # Process ['ListAvailable'] or logon request for ['FriendlyName'] or ['TenantId'] #
        ###################################################################################

        else {

            ###################################################
            # Ensure we have a valid CredentialsObject loaded #
            ###################################################

            # Load the CredentialsObject from CredentialsDataFile if not provided
            if (-not $CredentialsObject -and $CredentialsDataFile) {
                Write-Verbose "Loading credentials from CredentialsDataFile [$CredentialsDataFile]"
                $CredentialsObject = Import-PowerShellDataFile -Path $CredentialsDataFile -ErrorAction Stop
            }
            # Future feature: Load the CredentialsObject from CredentialsJsonFile if not provided
            elseif (-not $CredentialsObject -and $CredentialsJsonFile) {
                Write-Verbose "Loading credentials from CredentialsJsonFile [$CredentialsJsonFile]"
                $CredentialsObject = Get-Content -Path $CredentialsJsonFile | ConvertFrom-Json -ErrorAction Stop
            }
            # Throw error if not found a CredentialsObject by this point
            elseif (-not $CredentialsObject) {
                throw "Unable to find credentials."
            }

            ###################################################################
            # If ListAvailable, return credential sets from CredentialsObject #
            ###################################################################

            if ($ListAvailable) {
                $AvailableCredentials = Get-CredentialsList $CredentialsObject
                return $AvailableCredentials
            }

            ###############################################################################################
            # Get the CredentialsBlock using Get-CredentialsBlockByName or Get-CredentialsBlockByTenantId #
            ###############################################################################################

            if ($Name -and $TenantId) {
                Write-Verbose "Looking for credential block matching Name [$Name] and TenantId [$TenantId]"
                $CredentialsBlock = Get-CredentialsBlockByName -InputObject $CredentialsObject -Name $Name -TenantId $TenantId -ErrorAction Stop
            }
            elseif ($Name) {
                Write-Verbose "Looking for credential block matching Name [$Name]"
                $CredentialsBlock = Get-CredentialsBlockByName -InputObject $CredentialsObject -Name $Name -ErrorAction Stop
            }
            elseif ($TenantId) {
                Write-Verbose "Looking for credential block matching TenantId [$TenantId]"
                $CredentialsBlock = Get-CredentialsBlockByTenantId -InputObject $CredentialsObject -TenantId $TenantId -ErrorAction Stop
            }
            else {
                Write-Warning "Unable to process CredentialsBlock request"
            }

            #####################################################
            # Use the CredentialsBlock to set the Azure Context #
            #####################################################

            if ($CredentialsBlock) {
                if ($UseAzLogin) {
                    Write-Verbose "Switching Azure Context using Client ID [$($CredentialsBlock.clientId)] (az login)"
                    $ctx = az login --service-principal -u ($CredentialsBlock.clientId) -p ($CredentialsBlock.clientSecret) --tenant ($CredentialsBlock.tenantId) | ConvertFrom-Json
                    if (!$ctx) {
                        Write-Error "Error running az login"
                        return
                    }
                }
                else {
                    Write-Verbose "Switching Azure Context using Client ID [$($CredentialsBlock.clientId)]"
                    $Credential = New-Object System.Management.Automation.PSCredential (
                        $($CredentialsBlock.clientId),
                        $($CredentialsBlock.clientSecret | ConvertTo-SecureString -AsPlainText -Force)
                    )
                    Clear-AzContext -Force | Out-Null
                    $ctx = Connect-AzAccount -ServicePrincipal -Tenant $($CredentialsBlock.tenantId) -Credential $Credential
                }
                if ($UseTerraform) {
                    Write-Verbose "Setting environment variables for Terraform Azure Provider"
                    $env:ARM_CLIENT_ID = $CredentialsBlock.clientId
                    $env:ARM_CLIENT_SECRET = $CredentialsBlock.clientSecret
                    $env:ARM_TENANT_ID = $CredentialsBlock.tenantId
                    if ($CredentialsBlock.subscriptionId) {
                        Write-Verbose "Setting default Subscription for Terraform Azure Provider"
                        $env:ARM_SUBSCRIPTION_ID = $CredentialsBlock.subscriptionId
                    }
                }
            }
            else {
                Write-Warning "Nothing to process."
            }

        }

        ##################
        # Return context #
        ##################

        return $ctx

    }
}