Locoxella/JIT-Jumpbox-Manager-README.md

## JIT-Jumpbox-Manager-README.md

      
    Raw
  

              JIT-Jumpbox-Manager-README.md
            
          
    JIT Jumpbox Manager (PowerShell)

A production-ready PowerShell script to automate the lifecycle of an Azure Just-In-Time (JIT) jumpbox. It handles starting the VM, attaching a public IP, requesting JIT access, keeping the session alive, and performing a clean shutdown.
File


start-prod-jumpbox.ps1

Requirements


Azure CLI: az
PowerShell Modules:

Az.Accounts
Az.Security


Azure Subscription:

Microsoft Defender for Servers Plan 2 enabled.


RBAC Permissions for the authenticated principal:

Read VM and networking resources.
Read/write JIT network access policies.
Request JIT access (Microsoft.Security/locations/jitNetworkAccessPolicies/initiate/action).


Authentication

Before running the script, ensure you are authenticated with both Azure CLI and Azure PowerShell using an account that has the required RBAC permissions.
Azure CLI:
az login
Azure PowerShell:
Connect-AzAccount
Basic Usage

./start-prod-jumpbox.ps1 `
  -SubscriptionId "<subscription-id>" `
  -ResourceGroup "<resource-group>" `
  -VmName "<jumpbox-vm-name>" `
  -PublicIpName "<public-ip-resource-name>"
Optional Parameters

-Location "<azure-location>"

Note: The VM's location is auto-detected by default. This parameter serves as a fallback.

What it Does


Validates Prerequisites: Checks for az and required PowerShell modules.
Sets Azure Context: Aligns the subscription context for both Azure CLI and Az PowerShell.
Starts VM: Starts the specified jumpbox VM if it's not already running.
Attaches Public IP: Associates the specified Public IP resource with the VM's network interface.
Requests JIT Access: Initiates a JIT access request for SSH on port 22.
Provides Connection Details: Displays ready-to-use SSH command examples.
Maintains JIT Session: Actively monitors the JIT window and renews it before expiration to prevent disconnects.
Automates Cleanup: On exit (keypress), JIT failure, or expiration, it cleanly detaches the Public IP and deallocates the VM to save costs.

Connection Examples

Once the script confirms JIT access is ready, you can connect using your preferred SSH method.
Interactive Shell

ssh -i <private-key> azureuser@<public-ip>
Remote Command Execution

Run a single command through the jumpbox.
ssh -i <private-key> azureuser@<public-ip> "kubectl get nodes"
Dynamic SOCKS Proxy

Create a local SOCKS proxy to route traffic through the jumpbox.
ssh -i <private-key> -D 1080 -N azureuser@<public-ip>
Reverse Tunnel

Expose a local port to the jumpbox network.
ssh -i <private-key> -R 9443:localhost:443 azureuser@<public-ip>
Behavior Notes


Independent Runtimes: The VM's runtime and the JIT access window are managed independently. The VM can be running even if the JIT window has expired.
Keep-Alive & Safety: The script includes a keep-alive loop for the JIT session. If a renewal fails or the window expires, the script automatically proceeds to the cleanup phase.
Cleanup Delay: The script polls for a keypress to exit. By default, this check occurs every 30 seconds, so cleanup may be delayed by up to that amount of time after you press a key.

Troubleshooting

"Azure PowerShell is not authenticated"

Run Connect-AzAccount -Subscription <subscription-id> to log in and set the correct subscription context.
"access request is not a subset of policy"

This error means the script's JIT request (e.g., source IP, time) is not allowed by the JIT policy configured in Microsoft Defender for Cloud. The script attempts to read the policy and make a valid request, but if the policy was recently changed, it might fail. Verify the port 22 settings for the JIT policy in the Azure portal.
SSH connection times out even when the VM is running

This typically means the JIT access window has expired or the renewal failed. Re-run the script to request a new JIT session.
Security Recommendations


Use a dedicated, least-privileged user account and SSH keypair for the jumpbox.
Keep JIT access windows as short as is practical.
Regularly audit JIT access activity in the Microsoft Defender for Cloud dashboard.
Implement a key rotation policy and avoid sharing private keys.


## jit-jumpbox-manager.ps1
#!/usr/bin/env pwsh
<#
.SYNOPSIS
    Starts a jumpbox VM, attaches a public IP, and requests native JIT access.
.DESCRIPTION
    This script:
    1. Checks prerequisites (az CLI, PowerShell Az modules)
    2. Starts the target jumpbox VM
    3. Attaches the specified public IP resource
    4. Requests JIT access via Azure Defender for Servers (native JIT)
    5. Waits for user keypress (with JIT keep-alive)
    6. Detaches public IP and deallocates the VM on cleanup

    NOTE: Requires Defender for Servers Plan 2, Az.Accounts, and Az.Security.
          Native JIT automatically manages NSG rules.
#>

param(
    [Parameter(Mandatory=$true)]
    [string]$SubscriptionId,

    [Parameter(Mandatory=$true)]
    [string]$ResourceGroup,

    [Parameter(Mandatory=$true)]
    [string]$VmName,

    [Parameter(Mandatory=$true)]
    [string]$PublicIpName,

    [string]$Location = "eastus",  # Fallback location; VM location is auto-detected and preferred
    [bool]$EnableNsgSshFallback = $true,
    [int]$SshProbeTimeoutMs = 7000
)

$ErrorActionPreference = "Stop"

function Check-Command {
    param([string]$Command)
    $null = Get-Command $Command -ErrorAction SilentlyContinue
    return $?
}

function Check-PowerShellModule {
    param([string]$ModuleName)
    $module = Get-Module -Name $ModuleName -ListAvailable -ErrorAction SilentlyContinue
    return $null -ne $module
}

function Test-TcpPort {
    param(
        [Parameter(Mandatory=$true)]
        [string]$TargetHost,

        [Parameter(Mandatory=$true)]
        [int]$Port,

        [int]$TimeoutMs = 7000
    )

    $client = New-Object System.Net.Sockets.TcpClient
    try {
        $connectTask = $client.ConnectAsync($TargetHost, $Port)
        if (-not $connectTask.Wait($TimeoutMs)) {
            return $false
        }

        return $client.Connected
    }
    catch {
        return $false
    }
    finally {
        if ($null -ne $client) {
            $client.Dispose()
        }
    }
}

function Get-VmPowerState {
    param(
        [Parameter(Mandatory=$true)]
        [string]$ResourceGroup,

        [Parameter(Mandatory=$true)]
        [string]$VmName
    )

    $status = az vm get-instance-view `
        --resource-group $ResourceGroup `
        --name $VmName `
        --query "instanceView.statuses[?starts_with(code, 'PowerState/')].displayStatus" `
        -o tsv

    return "$status".Trim()
}

function Get-NicPublicIpId {
    param(
        [Parameter(Mandatory=$true)]
        [string]$NicId
    )

    $publicIpId = az network nic show --ids $NicId --query "ipConfigurations[0].publicIPAddress.id" -o tsv
    return "$publicIpId".Trim()
}

function Restore-NsgRuleState {
    param(
        [Parameter(Mandatory=$true)]
        [string]$NsgResourceGroup,

        [Parameter(Mandatory=$true)]
        [string]$NsgName,

        [Parameter(Mandatory=$true)]
        [hashtable]$RuleState
    )

    $restoreProtocol = if ([string]::IsNullOrEmpty("$($RuleState.Protocol)")) { "*" } else { "$($RuleState.Protocol)" }
    $restoreAccess = if ([string]::IsNullOrEmpty("$($RuleState.Access)")) { "Deny" } else { "$($RuleState.Access)" }

    $restoreOutput = $null
    if ($RuleState.SourceAddressPrefixes.Count -gt 0) {
        if ($restoreProtocol -eq "*") {
            $restoreOutput = az network nsg rule update `
                --resource-group $NsgResourceGroup `
                --nsg-name $NsgName `
                --name $RuleState.Name `
                --access $restoreAccess `
                --protocol '*' `
                --source-address-prefixes $RuleState.SourceAddressPrefixes 2>&1
        }
        else {
            $restoreOutput = az network nsg rule update `
                --resource-group $NsgResourceGroup `
                --nsg-name $NsgName `
                --name $RuleState.Name `
                --access $restoreAccess `
                --protocol "$restoreProtocol" `
                --source-address-prefixes $RuleState.SourceAddressPrefixes 2>&1
        }
    }
    else {
        $restoreSourcePrefix = if ([string]::IsNullOrEmpty("$($RuleState.SourceAddressPrefix)")) { "*" } else { "$($RuleState.SourceAddressPrefix)" }

        if ($restoreProtocol -eq "*" -and $restoreSourcePrefix -eq "*") {
            $restoreOutput = az network nsg rule update `
                --resource-group $NsgResourceGroup `
                --nsg-name $NsgName `
                --name $RuleState.Name `
                --access $restoreAccess `
                --protocol '*' `
                --source-address-prefixes '*' 2>&1
        }
        elseif ($restoreProtocol -eq "*") {
            $restoreOutput = az network nsg rule update `
                --resource-group $NsgResourceGroup `
                --nsg-name $NsgName `
                --name $RuleState.Name `
                --access $restoreAccess `
                --protocol '*' `
                --source-address-prefixes "$restoreSourcePrefix" 2>&1
        }
        elseif ($restoreSourcePrefix -eq "*") {
            $restoreOutput = az network nsg rule update `
                --resource-group $NsgResourceGroup `
                --nsg-name $NsgName `
                --name $RuleState.Name `
                --access $restoreAccess `
                --protocol "$restoreProtocol" `
                --source-address-prefixes '*' 2>&1
        }
        else {
            $restoreOutput = az network nsg rule update `
                --resource-group $NsgResourceGroup `
                --nsg-name $NsgName `
                --name $RuleState.Name `
                --access $restoreAccess `
                --protocol "$restoreProtocol" `
                --source-address-prefixes "$restoreSourcePrefix" 2>&1
        }
    }

    if ($LASTEXITCODE -eq 0) {
        return $true
    }

    Write-Host "  ⚠ Primary restore failed for '$($RuleState.Name)': $restoreOutput" -ForegroundColor Yellow
    Write-Host "  Attempting secure fallback restore (Deny/*/*)..." -ForegroundColor Yellow

    $fallbackOutput = az network nsg rule update `
        --resource-group $NsgResourceGroup `
        --nsg-name $NsgName `
        --name $RuleState.Name `
        --access Deny `
        --protocol '*' `
        --source-address-prefixes '*' 2>&1

    if ($LASTEXITCODE -eq 0) {
        Write-Host "  ✓ Fallback restore succeeded for '$($RuleState.Name)'" -ForegroundColor Green
        return $true
    }

    Write-Host "  ⚠ Fallback restore failed for '$($RuleState.Name)': $fallbackOutput" -ForegroundColor Red
    return $false
}

Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host "  Jumpbox VM Startup Script (Native JIT)" -ForegroundColor Cyan
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host ""

# Check prerequisites
Write-Host "[1/8] Checking prerequisites..." -ForegroundColor Yellow

$missingRequirements = @()

if (-not (Check-Command "az")) {
    $missingRequirements += "Azure CLI (az) - Install from: https://learn.microsoft.com/en-us/cli/azure/install-azure-cli"
}
else {
    Write-Host "  ✓ Azure CLI found" -ForegroundColor Green
}

if (-not (Check-PowerShellModule "Az.Security")) {
    $missingRequirements += "PowerShell Module 'Az.Security' - Install with: Install-Module -Name Az.Security -AllowClobber -Force"
}
else {
    Write-Host "  ✓ Az.Security PowerShell module found" -ForegroundColor Green
}

if (-not (Check-PowerShellModule "Az.Accounts")) {
    $missingRequirements += "PowerShell Module 'Az.Accounts' - Install with: Install-Module -Name Az.Accounts -AllowClobber -Force"
}
else {
    Write-Host "  ✓ Az.Accounts PowerShell module found" -ForegroundColor Green
}

if ($missingRequirements.Count -gt 0) {
    Write-Host ""
    Write-Host "  ⚠ Missing prerequisites:" -ForegroundColor Red
    foreach ($req in $missingRequirements) {
        Write-Host "  • $req" -ForegroundColor Red
    }
    exit 1
}

Write-Host "  ✓ All prerequisites met" -ForegroundColor Green
Write-Host ""
Write-Host "[2/8] Setting Azure subscription context and loading modules..." -ForegroundColor Yellow
az account set --subscription $SubscriptionId
if ($LASTEXITCODE -ne 0) {
    Write-Error "Failed to set subscription. Please ensure you're logged in with 'az login'"
    exit 1
}
Write-Host "  ✓ Subscription set: $SubscriptionId" -ForegroundColor Green

Import-Module Az.Accounts -ErrorAction SilentlyContinue
Import-Module Az.Security -ErrorAction Stop
Write-Host "  ✓ Az.Security module loaded" -ForegroundColor Green

# Ensure Azure PowerShell is authenticated and using the same subscription
$azPsContext = Get-AzContext -ErrorAction SilentlyContinue
if ($null -eq $azPsContext) {
    Write-Error "Azure PowerShell is not authenticated. Please run 'Connect-AzAccount -Subscription $SubscriptionId' and retry."
    exit 1
}

try {
    Set-AzContext -Subscription $SubscriptionId -ErrorAction Stop | Out-Null
    Write-Host "  ✓ Az PowerShell context set: $SubscriptionId" -ForegroundColor Green
}
catch {
    Write-Error "Failed to set Az PowerShell context to subscription '$SubscriptionId'. Error: $($_.Exception.Message)"
    exit 1
}
Write-Host ""

# Check current VM status
Write-Host "[3/8] Checking VM current state..." -ForegroundColor Yellow
$vmStatus = Get-VmPowerState -ResourceGroup $ResourceGroup -VmName $VmName
$vmStartedByScript = $false
$vmWasAlreadyRunning = $false
Write-Host "  Current VM Status: $vmStatus" -ForegroundColor Cyan

# Start the VM if not running
if ($vmStatus -notlike "*running*") {
    Write-Host "  Starting VM '$VmName'..." -ForegroundColor Yellow
    az vm start --resource-group $ResourceGroup --name $VmName --no-wait

    Write-Host "  Waiting for VM to start (this may take 1-2 minutes)..." -ForegroundColor Yellow
    $timeout = 180 # 3 minutes
    $elapsed = 0
    do {
        Start-Sleep -Seconds 5
        $elapsed += 5
        $vmStatus = Get-VmPowerState -ResourceGroup $ResourceGroup -VmName $VmName
        Write-Host "  Status: $vmStatus (${elapsed}s elapsed)" -ForegroundColor Gray

        if ($elapsed -ge $timeout) {
            Write-Error "VM failed to start within $timeout seconds"
            exit 1
        }
    } while ($vmStatus -notlike "*running*")

    $vmStartedByScript = $true
    Write-Host "  ✓ VM is now running (started by this script)" -ForegroundColor Green
} else {
    $vmWasAlreadyRunning = $true
    Write-Host "  ✓ VM is already running (not started by this script)" -ForegroundColor Green
}
Write-Host ""

# Get NIC information
Write-Host "[4/8] Getting VM network interface..." -ForegroundColor Yellow
$nicId = az vm show --resource-group $ResourceGroup --name $VmName --query "networkProfile.networkInterfaces[0].id" -o tsv
$nicName = ($nicId -split '/')[-1]
Write-Host "  NIC Name: $nicName" -ForegroundColor Cyan

# Get current IP configuration name
$ipConfigName = az network nic show --ids $nicId --query "ipConfigurations[0].name" -o tsv
$privateIpAddress = az network nic show --ids $nicId --query "ipConfigurations[0].privateIPAddress" -o tsv
$nsgId = az network nic show --ids $nicId --query "networkSecurityGroup.id" -o tsv
$nsgName = if ([string]::IsNullOrEmpty($nsgId)) { "" } else { ($nsgId -split '/')[-1] }
$nsgResourceGroup = if ([string]::IsNullOrEmpty($nsgId)) { $ResourceGroup } else { ($nsgId -split '/')[4] }
Write-Host "  IP Config Name: $ipConfigName" -ForegroundColor Cyan
Write-Host "  Private IP: $privateIpAddress" -ForegroundColor Cyan
if (-not [string]::IsNullOrEmpty($nsgName)) {
    Write-Host "  NSG Name: $nsgName" -ForegroundColor Cyan
}
Write-Host ""

# Check if public IP is already attached
Write-Host "[5/8] Checking public IP attachment..." -ForegroundColor Yellow
$currentPublicIp = Get-NicPublicIpId -NicId $nicId
$expectedPublicIpId = az network public-ip show --resource-group $ResourceGroup --name $PublicIpName --query "id" -o tsv

if ([string]::IsNullOrEmpty($currentPublicIp)) {
    Write-Host "  Attaching public IP '$PublicIpName' to NIC..." -ForegroundColor Yellow
    $null = az network nic ip-config update `
        --resource-group $ResourceGroup `
        --nic-name $nicName `
        --name $ipConfigName `
        --public-ip-address $PublicIpName 2>&1

    if ($LASTEXITCODE -ne 0) {
        Write-Error "Failed to attach public IP"
        exit 1
    }
    Write-Host "  ✓ Public IP attach requested" -ForegroundColor Green
} else {
    $currentIpName = ($currentPublicIp -split '/')[-1]
    if ($currentIpName -eq $PublicIpName) {
        Write-Host "  ✓ Public IP '$PublicIpName' is already attached" -ForegroundColor Green
    } else {
        Write-Host "  ⚠ Different public IP is attached: $currentIpName" -ForegroundColor Yellow
        Write-Host "  Replacing with '$PublicIpName'..." -ForegroundColor Yellow
        $null = az network nic ip-config update `
            --resource-group $ResourceGroup `
            --nic-name $nicName `
            --name $ipConfigName `
            --public-ip-address $PublicIpName 2>&1
        Write-Host "  ✓ Public IP replaced" -ForegroundColor Green
    }
}

# Validate IP attachment reached desired state (eventual consistency safe)
$attachTimeoutSeconds = 60
$attachElapsed = 0
do {
    $attachedPublicIp = Get-NicPublicIpId -NicId $nicId
    if ($attachedPublicIp -eq $expectedPublicIpId) {
        Write-Host "  ✓ Public IP is attached to NIC" -ForegroundColor Green
        break
    }

    Start-Sleep -Seconds 3
    $attachElapsed += 3
} while ($attachElapsed -lt $attachTimeoutSeconds)

if ($attachedPublicIp -ne $expectedPublicIpId) {
    Write-Error "Public IP '$PublicIpName' was not attached to NIC within $attachTimeoutSeconds seconds"
    exit 1
}

# Get the public IP address
$publicIpAddress = az network public-ip show --resource-group $ResourceGroup --name $PublicIpName --query "ipAddress" -o tsv
Write-Host "  Public IP Address: $publicIpAddress" -ForegroundColor Cyan
Write-Host ""

# Check NSG rules for SSH and request JIT access
Write-Host "[6/8] Requesting native JIT access for SSH..." -ForegroundColor Yellow

# Get current user's public IP
Write-Host "  Detecting your public IP address..." -ForegroundColor Gray
try {
    $userPublicIp = (Invoke-RestMethod -Uri "https://ipinfo.io/ip" -ErrorAction SilentlyContinue).Trim()
    if ([string]::IsNullOrEmpty($userPublicIp)) {
        $userPublicIp = (Invoke-RestMethod -Uri "https://api.ipify.org?format=json" -ErrorAction SilentlyContinue).ip
    }
    Write-Host "  Your public IP: $userPublicIp" -ForegroundColor Cyan
} catch {
    Write-Host "  ⚠ Could not detect your public IP automatically" -ForegroundColor Yellow
    $userPublicIp = Read-Host "Please enter your public IP address"
}

# Get the VM resource ID
$vmDetails = az vm show --resource-group $ResourceGroup --name $VmName -o json | ConvertFrom-Json
$vmResourceId = $vmDetails.id
$vmLocation = $vmDetails.location
$vmResourceGroupActual = ($vmResourceId -split '/')[4]

# Use VM-derived scope for JIT policy operations to avoid casing/scope mismatches
$jitResourceGroup = if ([string]::IsNullOrEmpty($vmResourceGroupActual)) { $ResourceGroup } else { $vmResourceGroupActual }
$jitLocation = if ([string]::IsNullOrEmpty($vmLocation)) { $Location } else { $vmLocation }
$jitPolicyResourceId = "/subscriptions/$SubscriptionId/resourceGroups/$jitResourceGroup/providers/Microsoft.Security/locations/$jitLocation/jitNetworkAccessPolicies/default"

Write-Host "  VM Resource ID: $vmResourceId" -ForegroundColor Cyan
Write-Host "  JIT Policy Scope RG: $jitResourceGroup" -ForegroundColor Cyan
Write-Host "  JIT Policy Scope Location: $jitLocation" -ForegroundColor Cyan

# Check if JIT policy exists for this VM
Write-Host "  Checking for existing JIT policies..." -ForegroundColor Gray
$jitPolicyObject = Get-AzJitNetworkAccessPolicy -ResourceGroupName $jitResourceGroup -Location $jitLocation -Name "default" -ErrorAction SilentlyContinue
$existingPolicies = $jitPolicyObject | Where-Object {
    $_.VirtualMachines | Where-Object { $_.Id -eq $vmResourceId }
}

if ($existingPolicies.Count -eq 0) {
    Write-Host "  No existing JIT policy found. Creating JIT policy..." -ForegroundColor Yellow

    # Create JIT policy with SSH port 22
    try {
        $jitPolicy = @{
            id = $vmResourceId
            ports = @(
                @{
                    number = 22
                    protocol = "*"
                    allowedSourceAddressPrefix = @("*")
                    maxRequestAccessDuration = "PT3H"  # 3 hours max
                }
            )
        }

        # Native JIT is managed through Az.Security PowerShell cmdlets.
        # It is not available as an Azure CLI (az) command yet.
        # TODO: Re-check Azure CLI support in future releases and simplify if native az support is added.
        Set-AzJitNetworkAccessPolicy -Kind "Basic" `
            -Location $jitLocation `
            -Name "default" `
            -ResourceGroupName $jitResourceGroup `
            -VirtualMachine @($jitPolicy) | Out-Null

        Write-Host "  ✓ JIT policy created successfully" -ForegroundColor Green
    } catch {
        Write-Host "  ⚠ Failed to create JIT policy: $($_.Exception.Message)" -ForegroundColor Yellow
        Write-Host "  Ensure Az PowerShell login: Connect-AzAccount -Subscription $SubscriptionId" -ForegroundColor Yellow
        Write-Host "  Continuing without JIT... SSH may be blocked" -ForegroundColor Yellow
    }
}
else {
    Write-Host "  ✓ JIT policy already exists for this VM" -ForegroundColor Green
}

# Request JIT access for up to 2 hours (bounded by policy max duration)
Write-Host "  Requesting JIT access for port 22 (up to 2 hours)..." -ForegroundColor Yellow

$requestDuration = [TimeSpan]::FromHours(2)
$requestedSourcePrefix = "$userPublicIp/32"
$policyAllowedPrefixString = ""

# Read policy constraints for this VM/port and keep request strictly within allowed subset
$policyVm = $null
$policyPort22 = $null
if ($null -ne $jitPolicyObject) {
    $policyVm = @($jitPolicyObject.VirtualMachines | Where-Object { $_.Id -eq $vmResourceId }) | Select-Object -First 1
    if ($null -ne $policyVm) {
        $policyPort22 = @($policyVm.Ports | Where-Object { $_.number -eq 22 }) | Select-Object -First 1
    }
}

if ($null -ne $policyPort22) {
    $allowedPrefixes = @()

    $policyAllowedPrefixString = "$($policyPort22.AllowedSourceAddressPrefix)".Trim()
    if (-not [string]::IsNullOrEmpty($policyAllowedPrefixString)) {
        $allowedPrefixes += ($policyAllowedPrefixString -split '\\s+' | Where-Object { -not [string]::IsNullOrWhiteSpace($_) })
    }

    if ($null -ne $policyPort22.AllowedSourceAddressPrefixes) {
        $allowedPrefixes += @($policyPort22.AllowedSourceAddressPrefixes)
    }

    $allowedPrefixes = @($allowedPrefixes | Select-Object -Unique)

    if (-not [string]::IsNullOrEmpty($policyPort22.maxRequestAccessDuration)) {
        try {
            $maxDuration = [System.Xml.XmlConvert]::ToTimeSpan($policyPort22.maxRequestAccessDuration)
            if ($requestDuration -gt $maxDuration) {
                $requestDuration = $maxDuration
                Write-Host "  ℹ Request duration adjusted to policy max: $($policyPort22.maxRequestAccessDuration)" -ForegroundColor Cyan
            }
        } catch {
            Write-Host "  ⚠ Could not parse policy maxRequestAccessDuration '$($policyPort22.maxRequestAccessDuration)'; using 2 hours" -ForegroundColor Yellow
        }
    }

    if ($allowedPrefixes.Count -gt 0 -and ($allowedPrefixes -notcontains "*") -and ($allowedPrefixes -notcontains "0.0.0.0/0")) {
        if ($allowedPrefixes -contains "$userPublicIp/32") {
            $requestedSourcePrefix = "$userPublicIp/32"
        }
        elseif ($allowedPrefixes -contains "$userPublicIp") {
            $requestedSourcePrefix = "$userPublicIp"
        }
        else {
            $requestedSourcePrefix = $allowedPrefixes[0]
            Write-Host "  ℹ Source prefix adjusted to policy-allowed value: $requestedSourcePrefix" -ForegroundColor Cyan
        }
    }
}

$endTimeUtc = (Get-Date).Add($requestDuration).ToUniversalTime().ToString("yyyy-MM-ddTHH:mm:ssZ")

$policyExists = Get-AzJitNetworkAccessPolicy -ResourceGroupName $jitResourceGroup -Location $jitLocation -Name "default" -ErrorAction SilentlyContinue
$jitRequestSucceeded = $false
$currentJitEndUtc = $null
if ($null -eq $policyExists) {
    Write-Host "  ⚠ JIT policy 'default' not found at scope: $jitPolicyResourceId" -ForegroundColor Yellow
    Write-Host "  Skipping JIT access request." -ForegroundColor Yellow
}
else {
    try {
        $requestPrefixesToTry = @($requestedSourcePrefix)
        if (-not [string]::IsNullOrEmpty($policyAllowedPrefixString) -and ($requestPrefixesToTry -notcontains $policyAllowedPrefixString)) {
            $requestPrefixesToTry += $policyAllowedPrefixString
        }

        $jitRequestSubmitted = $false
        foreach ($prefix in $requestPrefixesToTry) {
            try {
                $jitAccessRequest = @{
                    id = $vmResourceId
                    ports = @(
                        @{
                            number = 22
                            endTimeUtc = $endTimeUtc
                            allowedSourceAddressPrefix = $prefix
                        }
                    )
                }

                Start-AzJitNetworkAccessPolicy -ResourceId $jitPolicyResourceId `
                    -VirtualMachine @($jitAccessRequest) -ErrorAction Stop | Out-Null

                $requestedSourcePrefix = $prefix
                $jitRequestSubmitted = $true
                break
            }
            catch {
                Write-Host "  ⚠ JIT request submit failed for source prefix '$prefix': $($_.Exception.Message)" -ForegroundColor Yellow
            }
        }

        if (-not $jitRequestSubmitted) {
            throw "Failed to submit JIT request for all candidate source prefixes."
        }

        Write-Host "  ✓ JIT request submitted" -ForegroundColor Green
        Write-Host "  Requested policy source prefix: $requestedSourcePrefix" -ForegroundColor Cyan

        # JIT submission can succeed while the backend NSG update fails asynchronously.
        # Poll request status and only mark success when status is confirmed.
        $jitStatusTimeoutSeconds = 45
        $jitPollIntervalSeconds = 5
        $jitStatusDeadline = (Get-Date).ToUniversalTime().AddSeconds($jitStatusTimeoutSeconds)
        $jitRequestStatus = ""
        $jitRequestStatusReason = ""

        do {
            Start-Sleep -Seconds $jitPollIntervalSeconds

            $policyStatus = Get-AzJitNetworkAccessPolicy -ResourceGroupName $jitResourceGroup -Location $jitLocation -Name "default" -ErrorAction SilentlyContinue
            if ($null -eq $policyStatus -or $null -eq $policyStatus.Requests) {
                continue
            }

            $latestRequest = @($policyStatus.Requests | Sort-Object {
                try { [datetime]$_.StartTimeUtc } catch { [datetime]::MinValue }
            } -Descending) | Select-Object -First 1

            if ($null -eq $latestRequest) {
                continue
            }

            $latestVm = @($latestRequest.VirtualMachines | Where-Object { $_.Id -eq $vmResourceId }) | Select-Object -First 1
            if ($null -eq $latestVm) {
                continue
            }

            $latestPort = @($latestVm.Ports | Where-Object { $_.Number -eq 22 }) | Select-Object -First 1
            if ($null -eq $latestPort) {
                continue
            }

            $jitRequestStatus = "$($latestPort.Status)"
            $jitRequestStatusReason = "$($latestPort.StatusReason)"

            if ($jitRequestStatus -ieq "Failed") {
                $jitRequestSucceeded = $false
                break
            }

            if (
                $jitRequestStatus -ieq "Succeeded" -or
                $jitRequestStatus -ieq "Success" -or
                $jitRequestStatus -ieq "Approved" -or
                $jitRequestStatus -ieq "Active"
            ) {
                $jitRequestSucceeded = $true
                break
            }
        } while ((Get-Date).ToUniversalTime() -lt $jitStatusDeadline)

        if ($jitRequestSucceeded) {
            $currentJitEndUtc = [datetime]::Parse($endTimeUtc).ToUniversalTime()
            Write-Host "  ✓ JIT access granted successfully" -ForegroundColor Green
            Write-Host "  Access valid until: $endTimeUtc" -ForegroundColor Cyan
        }
        elseif (-not [string]::IsNullOrEmpty($jitRequestStatus)) {
            Write-Host "  ⚠ JIT request failed with status '$jitRequestStatus'" -ForegroundColor Yellow
            if (-not [string]::IsNullOrEmpty($jitRequestStatusReason)) {
                Write-Host "  Status reason: $jitRequestStatusReason" -ForegroundColor Yellow
            }
        }
        else {
            Write-Host "  ⚠ JIT request status was not confirmed within $jitStatusTimeoutSeconds seconds" -ForegroundColor Yellow
        }
    } catch {
        Write-Host "  ⚠ Failed to request JIT access: $($_.Exception.Message)" -ForegroundColor Yellow
        Write-Host "  Requested source prefix: $requestedSourcePrefix" -ForegroundColor Yellow
        Write-Host "  Ensure request values are a subset of policy for VM port 22." -ForegroundColor Yellow
        Write-Host "  Check Azure Portal → Defender for Cloud → Just-in-time VM access" -ForegroundColor Yellow
    }
}

# Validate SSH reachability and optionally apply NSG fallback for current private destination rules
$nsgFallbackRestoreRules = @()
$fallbackAllowSourcePrefix = "$userPublicIp/32"
$sshPortReachable = Test-TcpPort -TargetHost $publicIpAddress -Port 22 -TimeoutMs $SshProbeTimeoutMs
if ($sshPortReachable) {
    Write-Host "  ✓ SSH port 22 is reachable at $publicIpAddress" -ForegroundColor Green
}
elseif ($EnableNsgSshFallback -and -not [string]::IsNullOrEmpty($nsgName)) {
    Write-Host "  ⚠ SSH port 22 is not reachable. Testing NSG fallback by toggling Defender JIT deny rules..." -ForegroundColor Yellow
    $nsgRules = az network nsg rule list --resource-group $nsgResourceGroup --nsg-name $nsgName -o json | ConvertFrom-Json
    $candidateRules = @($nsgRules | Where-Object {
        $_.direction -eq "Inbound" -and
        $_.access -eq "Deny" -and
        (
            $_.destinationPortRange -eq "22" -or
            ($null -ne $_.destinationPortRanges -and $_.destinationPortRanges -contains "22")
        ) -and
        $_.name -like "MicrosoftDefenderForCloud-JITRule*"
    } | Sort-Object priority)

    if ($candidateRules.Count -eq 0) {
        Write-Host "  ⚠ No Defender JIT deny rules for port 22 were found in NSG '$nsgName'." -ForegroundColor Yellow
    }
    else {
        foreach ($rule in $candidateRules) {
            Write-Host "  Testing rule '$($rule.name)' (dest: $($rule.destinationAddressPrefix))..." -ForegroundColor Gray

            $originalSourcePrefixes = @()
            if ($null -ne $rule.sourceAddressPrefixes) {
                $originalSourcePrefixes = @($rule.sourceAddressPrefixes | Where-Object { -not [string]::IsNullOrWhiteSpace("$_") })
            }

            $originalRuleState = @{
                Name = $rule.name
                Access = if ([string]::IsNullOrEmpty("$($rule.access)")) { "Deny" } else { "$($rule.access)" }
                Protocol = if ([string]::IsNullOrEmpty("$($rule.protocol)")) { "*" } else { "$($rule.protocol)" }
                SourceAddressPrefix = "$($rule.sourceAddressPrefix)"
                SourceAddressPrefixes = $originalSourcePrefixes
            }

            $null = az network nsg rule update `
                --resource-group $nsgResourceGroup `
                --nsg-name $nsgName `
                --name $rule.name `
                --access Allow `
                --protocol Tcp `
                --source-address-prefixes $fallbackAllowSourcePrefix 2>&1

            if ($LASTEXITCODE -ne 0) {
                Write-Host "  ⚠ Failed to toggle rule '$($rule.name)' to Allow" -ForegroundColor Yellow
                continue
            }

            Start-Sleep -Seconds 3
            $sshPortReachable = Test-TcpPort -TargetHost $publicIpAddress -Port 22 -TimeoutMs $SshProbeTimeoutMs

            if ($sshPortReachable) {
                Write-Host "  ✓ SSH became reachable after toggling '$($rule.name)' to Allow" -ForegroundColor Green
                Write-Host "  ℹ Rule destination matched active jumpbox path. Rule will be restored to Deny during cleanup." -ForegroundColor Cyan
                $nsgFallbackRestoreRules += $originalRuleState
                break
            }
            else {
                $restored = Restore-NsgRuleState `
                    -NsgResourceGroup $nsgResourceGroup `
                    -NsgName $nsgName `
                    -RuleState $originalRuleState

                if ($restored) {
                    Write-Host "  Rule '$($rule.name)' is not the active path; restored to '$($originalRuleState.Access)'." -ForegroundColor Gray
                }
                else {
                    Write-Host "  ⚠ Rule '$($rule.name)' could not be restored automatically." -ForegroundColor Red
                }
            }
        }

        if (-not $sshPortReachable) {
            Write-Host "  ⚠ NSG fallback test did not restore SSH reachability." -ForegroundColor Yellow
        }
    }
}
elseif (-not $EnableNsgSshFallback) {
    Write-Host "  ⚠ SSH port 22 is not reachable and NSG fallback is disabled." -ForegroundColor Yellow
}
else {
    Write-Host "  ⚠ SSH port 22 is not reachable and NSG was not detected on the NIC." -ForegroundColor Yellow
}

Write-Host ""

# Display connection information
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host "  VM Ready for Connection" -ForegroundColor Green
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host "  Public IP:   $publicIpAddress" -ForegroundColor White
Write-Host "  SSH:         ssh -i <private-key> azureuser@$publicIpAddress" -ForegroundColor White
Write-Host ("  Command:     ssh -i <private-key> azureuser@{0} `"<command>`"" -f $publicIpAddress) -ForegroundColor White
Write-Host ("  Example:     ssh -i <private-key> azureuser@{0} `"kubectl config get-contexts -o name`"" -f $publicIpAddress) -ForegroundColor White
Write-Host "  Local SOCKS: ssh -i <private-key> -D 1080 -N azureuser@$publicIpAddress" -ForegroundColor White
Write-Host "  Reverse tun: ssh -i <private-key> -R 9443:localhost:443 azureuser@$publicIpAddress" -ForegroundColor White
Write-Host "  Note:        Replace <private-key> with your local key path" -ForegroundColor Gray
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host ""

# Wait for user input
Write-Host "[7/8] Press ANY KEY when you're done working with the VM..." -ForegroundColor Yellow
Write-Host "  (This will detach the public IP and stop the VM)" -ForegroundColor Yellow
$jitRenewBeforeExpiryMinutes = 10
$jitCheckIntervalSeconds = 30
Write-Host "  (Keypress is checked every $jitCheckIntervalSeconds seconds; cleanup may start up to that delay)" -ForegroundColor Yellow
$lastRenewAttemptUtc = [datetime]::MinValue
$autoStopDueToJit = -not $jitRequestSucceeded
$autoStopReason = ""

if ($autoStopDueToJit -and $sshPortReachable) {
    Write-Host "  ℹ JIT request did not succeed, but SSH is reachable via NSG fallback; keeping VM available." -ForegroundColor Cyan
    $autoStopDueToJit = $false
}

if ($autoStopDueToJit) {
    Write-Host "  ✖ JIT access was not granted. Jumpbox is not reachable from this machine." -ForegroundColor Red
    Write-Host "  Proceeding to immediate cleanup and VM deallocation." -ForegroundColor Red
}

if ($jitRequestSucceeded -and $null -ne $currentJitEndUtc) {
    Write-Host "  ℹ JIT keep-alive enabled (renews when <= $jitRenewBeforeExpiryMinutes minutes remain)" -ForegroundColor Cyan
}
else {
    Write-Host "  ℹ JIT keep-alive not active (initial JIT request did not succeed)" -ForegroundColor Yellow
}

while (-not $autoStopDueToJit) {
    if ($Host.UI.RawUI.KeyAvailable) {
        $null = $Host.UI.RawUI.ReadKey("NoEcho,IncludeKeyDown")
        break
    }

    $liveVmStatus = Get-VmPowerState -ResourceGroup $ResourceGroup -VmName $VmName
    if ($liveVmStatus -notlike "*running*") {
        Write-Host "  ⚠ VM state changed externally: $liveVmStatus" -ForegroundColor Yellow
        Write-Host "  Proceeding to cleanup..." -ForegroundColor Yellow
        $autoStopReason = "VM is no longer running ($liveVmStatus)"
        $autoStopDueToJit = $true
        break
    }

    $liveAttachedPublicIp = Get-NicPublicIpId -NicId $nicId
    if ($liveAttachedPublicIp -ne $expectedPublicIpId) {
        Write-Host "  ⚠ Public IP '$PublicIpName' is no longer attached to the jumpbox NIC." -ForegroundColor Yellow
        Write-Host "  Proceeding to cleanup..." -ForegroundColor Yellow
        $autoStopReason = "Public IP attachment changed externally"
        $autoStopDueToJit = $true
        break
    }

    if ($jitRequestSucceeded -and $null -ne $currentJitEndUtc) {
        $nowUtc = [datetime]::UtcNow
        $remaining = $currentJitEndUtc - $nowUtc

        if ($remaining.TotalSeconds -le 0) {
            Write-Host "  ⚠ JIT access window expired. Proceeding to cleanup..." -ForegroundColor Yellow
            $autoStopDueToJit = $true
            break
        }

        if ($remaining.TotalMinutes -le $jitRenewBeforeExpiryMinutes -and ($nowUtc - $lastRenewAttemptUtc).TotalSeconds -ge $jitCheckIntervalSeconds) {
            Write-Host "  ℹ Renewing JIT access before expiration..." -ForegroundColor Cyan
            $lastRenewAttemptUtc = $nowUtc

            try {
                $renewedEndTimeUtc = (Get-Date).Add($requestDuration).ToUniversalTime().ToString("yyyy-MM-ddTHH:mm:ssZ")
                $renewRequest = @{
                    id = $vmResourceId
                    ports = @(
                        @{
                            number = 22
                            endTimeUtc = $renewedEndTimeUtc
                            allowedSourceAddressPrefix = $requestedSourcePrefix
                        }
                    )
                }

                Start-AzJitNetworkAccessPolicy -ResourceId $jitPolicyResourceId `
                    -VirtualMachine @($renewRequest) | Out-Null

                $currentJitEndUtc = [datetime]::Parse($renewedEndTimeUtc).ToUniversalTime()
                Write-Host "  ✓ JIT access renewed until: $renewedEndTimeUtc" -ForegroundColor Green
            }
            catch {
                Write-Host "  ⚠ Failed to renew JIT access: $($_.Exception.Message)" -ForegroundColor Yellow
                Write-Host "  Proceeding to cleanup to avoid leaving VM running without access." -ForegroundColor Yellow
                $autoStopDueToJit = $true
                break
            }
        }
    }

    Start-Sleep -Seconds $jitCheckIntervalSeconds
}

if ($autoStopDueToJit) {
    if ([string]::IsNullOrEmpty($autoStopReason)) {
        Write-Host "  ℹ Cleanup is being triggered because JIT access is no longer available." -ForegroundColor Yellow
    }
    else {
        Write-Host "  ℹ Cleanup is being triggered because: $autoStopReason" -ForegroundColor Yellow
    }
}
Write-Host ""

# Cleanup: Detach public IP
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host "  Cleanup: Detaching Public IP and Stopping VM" -ForegroundColor Yellow
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host ""
Write-Host "[8/8] Performing cleanup..." -ForegroundColor Yellow

# Note: JIT access is automatically revoked by Azure after the time expires
# No manual NSG rule deletion needed - this is a major advantage of native JIT!
Write-Host "  ℹ JIT access will be automatically revoked when time expires" -ForegroundColor Cyan
Write-Host "  ℹ No manual NSG rule cleanup needed" -ForegroundColor Cyan
Write-Host ""

if ($nsgFallbackRestoreRules.Count -gt 0) {
    Write-Host "  Restoring temporary NSG fallback rules to Deny..." -ForegroundColor Yellow
    foreach ($ruleState in $nsgFallbackRestoreRules) {
        $restored = Restore-NsgRuleState `
            -NsgResourceGroup $nsgResourceGroup `
            -NsgName $nsgName `
            -RuleState $ruleState

        if ($restored) {
            Write-Host "  ✓ Restored '$($ruleState.Name)' to '$($ruleState.Access)'" -ForegroundColor Green
        }
        else {
            Write-Host "  ⚠ Failed to restore '$($ruleState.Name)'" -ForegroundColor Red
        }
    }
    Write-Host ""
}

Write-Host "  Detaching public IP from NIC..." -ForegroundColor Yellow
$publicIpBeforeDetach = Get-NicPublicIpId -NicId $nicId
if ([string]::IsNullOrEmpty($publicIpBeforeDetach)) {
    Write-Host "  ✓ Public IP is already detached (possibly external cleanup)" -ForegroundColor Green
}
elseif ($publicIpBeforeDetach -ne $expectedPublicIpId) {
    Write-Host "  ⚠ A different public IP is attached; skipping detach to avoid removing unexpected resource" -ForegroundColor Yellow
}
else {
    $null = az network nic ip-config update `
        --resource-group $ResourceGroup `
        --nic-name $nicName `
        --name $ipConfigName `
        --remove publicIpAddress 2>&1

    if ($LASTEXITCODE -eq 0) {
        Write-Host "  ✓ Public IP detached by this script" -ForegroundColor Green
    }
    else {
        Write-Host "  ⚠ Failed to detach public IP" -ForegroundColor Red
    }
}

# Stop the VM
$vmStatusBeforeStop = Get-VmPowerState -ResourceGroup $ResourceGroup -VmName $VmName
if ($vmStatusBeforeStop -like "*running*") {
    Write-Host "  Stopping VM '$VmName'..." -ForegroundColor Yellow
    $null = az vm deallocate --resource-group $ResourceGroup --name $VmName --no-wait 2>&1

    if ($LASTEXITCODE -eq 0) {
        if ($vmStartedByScript) {
            Write-Host "  ✓ VM stop initiated by this script (VM was started by this script)" -ForegroundColor Green
        }
        elseif ($vmWasAlreadyRunning) {
            Write-Host "  ✓ VM stop initiated by this script (VM was already running before script start)" -ForegroundColor Green
        }
        else {
            Write-Host "  ✓ VM stop initiated by this script" -ForegroundColor Green
        }
    }
    else {
        Write-Host "  ⚠ Failed to initiate VM stop" -ForegroundColor Red
    }
}
else {
    Write-Host "  ✓ VM is already stopped ($vmStatusBeforeStop)" -ForegroundColor Green
}

Write-Host ""
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host "  Cleanup Complete" -ForegroundColor Green
Write-Host "===========================================================" -ForegroundColor Cyan
Write-Host "  The VM will be fully stopped in a few minutes." -ForegroundColor Gray
Write-Host "  JIT access automatically expires after 2 hours." -ForegroundColor Gray
Write-Host ""
No results found