initiateAzureVmJit.ps1

# defines Azure VM JIT network access policies and uses own IP address to open up either SSH or RDP

param(
   # RegEx pattern to find your VM in your current subscription
   [Parameter(Mandatory = $true, Position = 1)]
   [string]$NamePattern,
   # uses ZScaler IP information to determine outbound IP to be opened up
   [switch]$AllowOffice = $false,
   [switch]$AllowVpn = $false,
   # fall-back to open up VM completely (NOT RECOMMENDED!)
   [switch]$AllowAll = $false,
   # show current requests
   [switch]$ShowRequests = $false
)

# assumptions:
# - up-to-date Azure PowerShell module installed
# - additionally module Az.Security installed
# - logged in and correct Azure subscription already selected

$VMs = Get-AzVM | Where-Object { $_.Name -match $NamePattern }

if ($VMs) {

   foreach ($vm in $VMs) {

      Write-Host "VM:" $vm.Name

      if ($vm.OSProfile.WindowsConfiguration) {
         $port = 3389
      }
      else {
         $port = 22
      }

      # set access policy
      $jitPolicy = (@{ id = $vm.Id; ports = (@{ number = $port; protocol = "*"; allowedSourceAddressPrefix = @("*"); maxRequestAccessDuration = "PT3H" }) })
                
      $accessPolicy = Set-AzJitNetworkAccessPolicy -Kind "Basic" -Location $vm.Location -Name $vm.Name -ResourceGroupName $vm.ResourceGroupName -VirtualMachine @($jitPolicy)

      $endTimeUtc = [DateTime]::UtcNow.AddHours(3).ToString('o')

      # start request
      # - if working from Office or Vpn try to figure out Zscaler outbound IP for RDP/SSH (my IP address is not accurate in this scenario)
      # - otherwise rely on my IP address
      if ($AllowOffice -or $AllowVpn) {
         if ($AllowOffice) {
            $searchPattern = "(?<=Your request is arriving at this server from the IP address.*\>)(\d+\.\d+\.\d+)"
         }
         else {
            $searchPattern = "(?<=The request is being received by the Zscaler Proxy from the IP address.*\>)(\d+\.\d+\.\d+)"
         }
         $zscalerContent = (Invoke-WebRequest -Uri "https://ip.zscaler.com/").Content.Split("`n")
         $zscalerIpRange = $zscalerContent | Select-String -Pattern $searchPattern -AllMatches | ForEach-Object { $_.Matches.Value }
         $myIp = @($zscalerIpRange + ".0/24")
      }
      elseif ($AllowAll) {
         $myIp = @("*")
      }
      else {
         $myIp = @((Invoke-RestMethod -Method get -Uri 'https://api.ipify.org?format=json').ip)
      }

      Write-Host "requesting access for IP" $myIp

      $jitPolicyVm = (@{ id = $vm.Id; ports = (@{ number = $port; endTimeUtc = $endTimeUtc; allowedSourceAddressPrefix = $myIp }) })

      $request = Start-AzJitNetworkAccessPolicy -ResourceId $accessPolicy.Id -VirtualMachine @($jitPolicyVm)

      # show all requests for debugging
      if ($ShowRequests) {
         Get-AzJitNetworkAccessPolicy -ResourceGroupName $vm.ResourceGroupName -Name $vm.Name -Location $vm.Location | `
            Select-Object -ExpandProperty Requests | `
            Select-Object -ExpandProperty VirtualMachines | `
            Select-Object -ExpandProperty Ports | `
            ? { $_.EndTimeUtc -gt $(Get-Date -AsUTC) }
      }

      # start VM to give request time to "settle"
      $tsBefore = Get-Date
      $vm | Start-AzVM
      $tsAfter = Get-Date
      $tsp = New-TimeSpan -Start $tsBefore -End $tsAfter
      if ($tsp.Seconds -lt 90) {
         Start-Sleep -Seconds $(90 - $tsp.Seconds)
      }

      # RDP connect for Windows / SSH for Linux
      $ipConfig = Get-AzNetworkInterface -ResourceId $vm.NetworkProfile.NetworkInterfaces[0].Id | Get-AzNetworkInterfaceIpConfig -ErrorAction Stop

      if ($ipConfig.PublicIpAddress) {
         $publicIp = Get-AzResource -ResourceId $ipConfig.PublicIpAddress.Id | Get-AzPublicIpAddress -ErrorAction Stop
         if ($publicIp.IpAddress) {
            if ($publicIp.DnsSettings.Fqdn) {
               $address = $publicIp.DnsSettings.Fqdn
            }
            else {
               $address = $publicIp.IpAddress
            }

            # WINDOWS
            if ($vm.OSProfile.WindowsConfiguration) {
               # if RDP file (with saved settings) exists, use it
               $rdpFileName = $(Join-Path $PSScriptRoot $($vm.Name + ".rdp"))
               if (Test-Path $rdpFileName) {
                  & $rdpFileName
               }
               # otherwise mstsc / RDP
               else {
                  & mstsc.exe /v:$address
               }
            }
            # LINUX
            else {
               $connectionString = $vm.OSProfile.AdminUsername + "@" + $address
               & ssh $connectionString
            }
         }
         else {
            Write-Error $("No public IP address information for " + $vm.Name)
         }
      }
      else {
         Write-Error $("No public IP address for " + $vm.Name)
      }
   }
}
else {
   Write-Error $("no VM with name pattern " + $NamePattern + " found!")
}