Add-JitAccess.ps1

#requires -Modules Az.Accounts
using namespace Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters

#function Add-JitAccess {
    <#
        .Synopsis
            Adds a JIT exception to VMs in the provided Resource Group
        .Notes
            Until we have an Az.Accounts version that works with Az.Security, we cannot use the provided cmdlets to maintain JIT. This is a workaround.
        .Example
            Add-JitAccess -SubscriptionName Subscription -ResourceGroupName LabGroup -VmName host
    #>
    param(
        # Subscription in which to find the JIT policies
        #[SubscriptionNameCompleter()]
        [string]$SubscriptionName,

        # Resource Group in which the VMs reside
        [Parameter(Mandatory)]
        [ResourceGroupCompleter()]
        [string]$ResourceGroupName,

        # VM to request JIT access to
        [ResourceNameCompleter("Microsoft.Compute/virtualMachines", "ResourceGroupName")]
        [string]$VMName = ".*",

        # Ports to request access to
        [uint16[]]$Port = 3389
    )
    begin {
        if ($SubscriptionName) {
            $null = Select-AzSubscription -SubscriptionName $SubscriptionName
        }

        $AzContext = Get-AzContext
        $AccessToken = if (Get-Command Get-AzAccessToken -ErrorAction SilentlyContinue) {
            (Get-AzAccessToken).Token  # Available in Az.Accounts >= 2.1
        } else {
            [Microsoft.Azure.Commands.Common.Authentication.AzureSession]::Instance.AuthenticationFactory.Authenticate(
                $AzContext.Account,
                $AzContext.Environment,
                $AzContext.Tenant.Id,
                $null,
                [Microsoft.Azure.Commands.Common.Authentication.ShowDialog]::Never,
                $null,
                'https://management.azure.com/'
            ).AccessToken
        }

        $Authentication = if ($PSVersionTable.PSVersion.Major -eq 7) {
            @{
                Authentication = "Bearer"
                Token          = ConvertTo-SecureString -String $AccessToken -AsPlainText -Force
            }
        } else {
            @{
                Headers = @{
                    Authorization = "Bearer $AccessToken"
                }
            }
        }

        $SubscriptionId = $AzContext.Subscription.Id
    }
    process {
        # Attempt to list policies
        $GetPolicies = @{
            Method         = "Get"
            Uri            = "https://management.azure.com/subscriptions/$SubscriptionId/providers/Microsoft.Security/jitNetworkAccessPolicies?api-version=2020-01-01"
            ContentType    = "application/json"
        }
        $JitPolicy = (Invoke-RestMethod @GetPolicies @Authentication).value.Where{$_.id -match "/resourceGroups/$($ResourceGroupName)/"}

        # Danger: Add handling for multiple policies? All current are "default", but that may change

        $JitRequest = @{
            Method         = "Post"
            Uri            = "https://management.azure.com/subscriptions/$($SubscriptionId)/resourceGroups/$($ResourceGroupName)/providers/Microsoft.Security/locations/$($JitPolicy.location)/jitNetworkAccessPolicies/$($JitPolicy.name)/initiate?api-version=2020-01-01"
            ContentType    = "application/json"
            Body           = @{
                virtualMachines = @(
                    foreach ($Vm in $JitPolicy.properties.virtualMachines.Where{$_.Id -match "/$($VMName)$"}) {
                        @{
                            id    = $Vm.Id
                            ports = @(
                                foreach ($i in $Port) {
                                    @{
                                        # allowedSourceAddressPrefix = # Defaults to source IP address of the request
                                        # endTimeUtc = Get-Date (Get-Date).AddHours(8).ToUniversalTime() -Format "o"
                                        duration = "PT8H"
                                        number   = $i
                                    }
                                }
                            )
                        }
                    }
                )
                justification   = "Add-QmLabJitAccess:$($env:UserName)"
            } | ConvertTo-Json -Depth 10
        }

        $JitInit = Invoke-RestMethod @JitRequest @Authentication

        Write-Verbose "JIT Connection to '$($JitInit.virtualMachines.id.Split('/')[-1]):$($JitInit.virtualMachines.ports.number)' $($JitInit.virtualMachines.ports[0].status)."
        Write-Verbose "Connections will be blocked after $($JitInit.virtualMachines.ports[0].endTimeUtc) UTC"
    }
#}