# Add Subscription IDs here that will be iterated
$SubscriptionIds = @(
    "Subscriptionid1",
    "Subscriptionid2"
)

# Add any virtual machines here you want to be excluded from JIT
$ExcludedVMs = @(
    "VMName1",
    "VMName2"
)

foreach ($SubscriptionId in $SubscriptionIds) {
    Write-Output "Attempting to set context for subscription: $SubscriptionId"

    try {
        Set-AzContext -SubscriptionId $SubscriptionId -ErrorAction Stop
        Write-Output "Processing subscription: $SubscriptionId"
    }
    catch {
        Write-Output "Failed to set context for subscription: $SubscriptionId"
        Write-Output $_
        continue
    }
    try {
        $ResourceGroups = Get-AzResourceGroup

        foreach ($ResourceGroup in $ResourceGroups) {
            $ResourceGroupName = $ResourceGroup.ResourceGroupName

            # Get VMs in the resource group
            $VMs = Get-AzVM -ResourceGroupName $ResourceGroupName

            # Only write output if there are VMs in the RG
            if ($VMs.Count -gt 0) {
                Write-Output "Processing resource group: $ResourceGroupName"
            }

            foreach ($VM in $VMs) {
                $VMName = $VM.Name

                # Check if the VM is in the exclusion list
                if ($ExcludedVMs -contains $VMName) {
                    Write-Output "Skipping VM (excluded): $VMName"
                    continue
                }

                # Check if JIT is already enabled for the VM
                $JitPolicies = Get-AzJitNetworkAccessPolicy -ResourceGroupName $ResourceGroupName
                $JitPolicyExists = $false

                foreach ($JitPolicy in $JitPolicies.VirtualMachines) {
                    if ($JitPolicy.id -eq "/subscriptions/$SubscriptionId/resourceGroups/$ResourceGroupName/providers/Microsoft.Compute/virtualMachines/$VMName") {
                        $JitPolicyExists = $true
                        break
                    }
                }

                if ($JitPolicyExists) {
                    Write-Output "JIT already enabled for VM: $VMName"
                } else {
                    # JIT policy configuration
                    $JitPolicyConfig = @(
                        @{
                            id = "/subscriptions/$SubscriptionId/resourceGroups/$ResourceGroupName/providers/Microsoft.Compute/virtualMachines/$VMName";
                            ports = @(
                                @{
                                    number = 22; # SSH port
                                    protocol = "*";
                                    allowedSourceAddressPrefix = @(""); # Add allowed source address prefix, Azure Bastion subnet
                                    maxRequestAccessDuration = "PT3H" # Set time window of 3 Hours. Maximum is 24 hours.
                                },
                                @{
                                    number = 3389; # RDP port
                                    protocol = "*";
                                    allowedSourceAddressPrefix = @(""); # Add allowed source address prefix, Azure Bastion subnet
                                    maxRequestAccessDuration = "PT3H" # Set time window of 3 hours. Maximum is 24 hours.
                                }
                            )
                        }
                    )

                    # Enable JIT for VM
                    Set-AzJitNetworkAccessPolicy -Kind "Basic" -Location $VM.Location -Name "${VMName}-JITPolicy" -ResourceGroupName $ResourceGroupName -VirtualMachine $JitPolicyConfig

                    Write-Output "JIT enabled for VM: $VMName"
                }
            }
        }
    }
    catch {
        Write-Output "An error occurred while processing subscription: $SubscriptionId"
        Write-Output $_
    }
}