#requires -Version 7.0
Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

function Initialize-Context($Root,$Subscription,$Region,$Size,$Timeout) {
    $script:Bundle = [IO.Path]::GetFullPath($Root)
    $script:Runtime = Join-Path (Split-Path (Split-Path $script:Bundle)) 'work/azure-monitor-private-link-deep-dive'
    $script:StateFile = Join-Path $script:Runtime 'state.json'
    $script:State = if (Test-Path $script:StateFile) { Get-Content $script:StateFile -Raw | ConvertFrom-Json -AsHashtable } else { $null }
    if ($script:State) { $script:State.Resources=[Collections.Generic.List[string]]@($script:State.Resources) }
    $script:RequestedSubscription = $Subscription
    $script:Location = $Region
    $script:VmSize = $Size
    $script:Timeout = $Timeout
    $script:Rg = 'rg-ampls-deepdive-we'
    $script:Api = @{Network='2024-05-01';Compute='2024-07-01';Insights='2023-03-11';Dcr='2023-03-11';Ampls='2023-06-01-preview';Logs='2023-09-01';Dns='2020-06-01';Identity='2023-01-31'}
    $script:Zones = @('privatelink.monitor.azure.com','privatelink.oms.opinsights.azure.com','privatelink.ods.opinsights.azure.com','privatelink.agentsvc.azure-automation.net','privatelink.blob.core.windows.net')
}
function Write-PrivateJson($Name,$Object) {
    [IO.File]::WriteAllText((Join-Path $script:Runtime $Name), ($Object | ConvertTo-Json -Depth 90), [Text.UTF8Encoding]::new($false))
}
function Save-State { Write-PrivateJson 'state.json' $script:State }
function Invoke-Az([string[]]$Arguments,[switch]$MissingOkay) {
    $oldWarnings = $env:PYTHONWARNINGS
    try {
        $env:PYTHONWARNINGS='ignore'
        $result = & (Get-Command az).Source @Arguments --only-show-errors --output json 2>&1
        $code = $LASTEXITCODE
    } finally { $env:PYTHONWARNINGS=$oldWarnings }
    $raw = $result -join "`n"
    if ($code -ne 0) {
        if ($MissingOkay -and $raw -match 'ResourceNotFound|ResourceGroupNotFound|NotFound|\bNot Found\b|could not be found') { return $null }
        if (Test-Path $script:Runtime) { Write-PrivateJson 'last-cli-error.json' @{Arguments=$Arguments;Output=$raw;ExitCode=$code} }
        throw "Azure CLI returned $code. $raw"
    }
    if ($raw.Trim()) { return ($raw | ConvertFrom-Json -AsHashtable -Depth 90) }
}
function Invoke-Arm($Method,$Id,$Version,$Body=$null,[switch]$MissingOkay) {
    $argsList = @('rest','--method',$Method,'--url',"https://management.azure.com${Id}?api-version=$Version")
    if ($null -ne $Body) {
        Write-PrivateJson 'request.json' $Body
        $argsList += @('--body',"@$(Join-Path $script:Runtime 'request.json')",'--headers','Content-Type=application/json')
    }
    for ($attempt=0; $attempt -lt 6; $attempt++) {
        try { return Invoke-Az $argsList -MissingOkay:$MissingOkay }
        catch {
            if ($attempt -eq 5 -or $_ -notmatch '429|TooManyRequests|InternalServerError|ServiceUnavailable|AnotherOperationInProgress|RetryableError|ConnectTimeout|ReadTimeout|ConnectionError') { throw }
            Start-Sleep -Seconds ([Math]::Min(30,[Math]::Pow(2,$attempt+1)))
        }
    }
}
function Resource-Id($Provider,$Name) { "$($script:State.GroupId)/providers/$Provider/$Name" }
function Assert-Ownership {
    if (-not $script:State) { throw 'No runtime ownership state. Run Foundation first.' }
    $account = Invoke-Az @('account','show')
    if ($account.id -ne $script:State.Subscription -or ($script:RequestedSubscription -and $script:RequestedSubscription -ne $account.id)) { throw 'Subscription mismatch. Select the recorded subscription explicitly with az account set.' }
    if ($script:Location -ne $script:State.Location -or $script:VmSize -ne $script:State.VmSize) { throw 'Location or VM size differs from recorded session.' }
    $group = Invoke-Az @('group','show','--name',$script:Rg)
    if ($group.tags['workshop'] -ne 'ampls-deepdive' -or $group.tags['session'] -ne $script:State.Session) { throw 'Unexpected resource-group ownership. Refusing mutation.' }
}
function Wait-Resource($Id,$Version) {
    $deadline=[datetime]::UtcNow.AddMinutes(20)
    do {
        $r=Invoke-Arm GET $Id $Version -MissingOkay
        if (-not $r) { Start-Sleep 10; continue }
        $status=$r.properties['provisioningState']
        if (-not $status -or $status -eq 'Succeeded') { return $r }
        if ($status -in @('Failed','Canceled')) { throw "Resource provisioning failed: $Id" }
        Start-Sleep 10
    } while ([datetime]::UtcNow -lt $deadline)
    throw "Resource provisioning timeout: $Id"
}
function Ensure-Resource($Id,$Version,$Body,[switch]$Child) {
    $old=Invoke-Arm GET $Id $Version -MissingOkay
    if ($old) {
        if (-not $script:State.Resources.Contains($Id)) { throw "Unexpected existing resource: $Id" }
        if (-not $Child -and $old.tags['session'] -ne $script:State.Session) { throw 'Resource ownership tag mismatch.' }
        return Wait-Resource $Id $Version
    }
    if (-not $script:State.Resources.Contains($Id)) { $script:State.Resources.Add($Id); Save-State }
    $null=Invoke-Arm PUT $Id $Version $Body
    Wait-Resource $Id $Version
}
function New-Body($Properties,$Region=$script:Location) { @{location=$Region;tags=$script:State.Tags;properties=$Properties} }
function Require-Stage($Name) { if (-not $script:State.Completed[$Name]) { throw "Required stage $Name has not passed." } }
function Invoke-Preflight {
    $account=Invoke-Az @('account','show')
    if ($account.state -ne 'Enabled') { throw 'Active subscription is not enabled.' }
    if ($script:RequestedSubscription -and $script:RequestedSubscription -ne $account.id) { throw 'Requested subscription is not active; run az account set explicitly, then retry.' }
    $null=Invoke-Az @('account','get-access-token','--resource','https://management.azure.com/')
    $exists=Invoke-Az @('group','exists','--name',$script:Rg)
    if ($exists) { Assert-Ownership }
    $checks=@{
        'Microsoft.Network'=@{virtualNetworks=$script:Api.Network;privateEndpoints=$script:Api.Network;natGateways=$script:Api.Network;publicIPAddresses=$script:Api.Network;privateDnsZones=$script:Api.Dns}
        'Microsoft.Compute'=@{virtualMachines=$script:Api.Compute}
        'Microsoft.ManagedIdentity'=@{userAssignedIdentities=$script:Api.Identity}
        'Microsoft.OperationalInsights'=@{workspaces=$script:Api.Logs}
        'Microsoft.Insights'=@{dataCollectionEndpoints=$script:Api.Insights;dataCollectionRules=$script:Api.Dcr;privateLinkScopes=$script:Api.Ampls}
    }
    foreach ($ns in $checks.Keys) {
        $p=Invoke-Az @('provider','show','--namespace',$ns)
        if ($p.registrationState -ne 'Registered') { throw "Provider $ns must already be registered. Preflight will not register it." }
        foreach ($type in $checks[$ns].Keys) {
            $t=@($p.resourceTypes | Where-Object resourceType -eq $type)
            if ($t.Count -ne 1 -or $checks[$ns][$type] -notin $t[0].apiVersions) { throw "Required API unavailable: $ns/$type" }
        }
    }
    $skus=@(Invoke-Az @('vm','list-skus','--location',$script:Location,'--size',$script:VmSize,'--all') | Where-Object name -eq $script:VmSize)
    if ($skus.Count -ne 1 -or @($skus[0].restrictions | Where-Object type -eq 'Location').Count) { throw 'VM size unavailable for the subscription and location.' }
    $cores=[int](@($skus[0].capabilities | Where-Object name -eq 'vCPUs')[0].value)
    $usage=Invoke-Az @('vm','list-usage','--location',$script:Location)
    foreach ($name in @('cores',$skus[0].family)) {
        $u=@($usage | Where-Object { $_.name.value -eq $name })
        if ($u.Count -ne 1 -or ([int]$u[0].limit-[int]$u[0].currentValue) -lt (2*$cores) -and -not $exists) { throw "Insufficient/unknown compute quota: $name" }
    }
    $networkUsage=Invoke-Az @('network','list-usages','--location',$script:Location)
    $pip=@($networkUsage | Where-Object { $_.name.value -eq 'PublicIPAddresses' })
    if ($pip.Count -ne 1 -or (([int]$pip[0].limit-[int]$pip[0].currentValue) -lt 2 -and -not $exists)) { throw 'Insufficient/unknown regional public-IP quota.' }
    $null=Invoke-Az @('vm','image','show','--location',$script:Location,'--urn','Canonical:ubuntu-24_04-lts:server:latest')
    $permissions=Invoke-Arm GET "/subscriptions/$($account.id)/providers/Microsoft.Authorization/permissions" '2022-04-01'
    if (-not @($permissions.value | Where-Object { '*' -in $_.actions -and @($_.notActions).Count -eq 0 }).Count) { throw 'This lab preflight requires unrestricted subscription control-plane permission (Owner); it will not guess effective custom permissions.' }
    Write-Host 'PASS: read-only authentication, ownership, API/provider, image, VM-size, quota and Owner permission checks.'
}
function Invoke-Foundation {
    Invoke-Preflight
    if ($script:State -and $script:State.Completed['Foundation']) { Assert-Ownership; return }
    if (-not $script:State) {
        $null=New-Item -ItemType Directory -Path $script:Runtime -Force
        if ($IsWindows) {
            $acl=Get-Acl $script:Runtime
            $acl.SetAccessRuleProtection($true,$false)
            $rule=[Security.AccessControl.FileSystemAccessRule]::new([Security.Principal.WindowsIdentity]::GetCurrent().User,'FullControl','ContainerInherit,ObjectInherit','None','Allow')
            $acl.SetAccessRule($rule); Set-Acl $script:Runtime $acl
        } else { & chmod 700 $script:Runtime }
        $account=Invoke-Az @('account','show')
        $session='ampls-'+[Convert]::ToHexString([Security.Cryptography.RandomNumberGenerator]::GetBytes(5)).ToLowerInvariant()
        $script:State=@{Session=$session;Subscription=$account.id;Location=$script:Location;VmSize=$script:VmSize;GroupId="/subscriptions/$($account.id)/resourceGroups/$script:Rg";Tags=@{workshop='ampls-deepdive';session=$session;environment='disposable-lab'};Completed=@{};Resources=[Collections.Generic.List[string]]::new();Evidence=@{};CreatedUtc=[datetime]::UtcNow.ToString('o');Locked=$false}
        Save-State
    }
    $script:State.Resources=[Collections.Generic.List[string]]@($script:State.Resources)
    if (-not (Invoke-Az @('group','exists','--name',$script:Rg))) {
        $null=Invoke-Az @('group','create','--name',$script:Rg,'--location',$script:Location,'--tags','workshop=ampls-deepdive',"session=$($script:State.Session)",'environment=disposable-lab')
    }
    Assert-Ownership
    $key=Join-Path $script:Runtime 'vm-key'
    if (-not (Test-Path "$key.pub")) {
        & ssh-keygen -t rsa -b 3072 -N '""' -f $key -q
        if ($LASTEXITCODE) { throw 'SSH key generation failed.' }
    }
    $publicKey=(Get-Content "$key.pub" -Raw).Trim()
    $uamiId=Resource-Id 'Microsoft.ManagedIdentity/userAssignedIdentities' 'id-ampls-probe'
    $identity=Ensure-Resource $uamiId $script:Api.Identity (New-Body @{})
    $script:State.Identity=$identity; Save-State
    foreach ($side in @('inside','outside')) {
        $prefix=if ($side -eq 'inside') {'10.94'} else {'10.95'}
        $pipId=Resource-Id 'Microsoft.Network/publicIPAddresses' "pip-ampls-$side"
        $pipBody=New-Body @{publicIPAllocationMethod='Static';publicIPAddressVersion='IPv4'}
        $pipBody.sku=@{name='Standard';tier='Regional'}
        $null=Ensure-Resource $pipId $script:Api.Network $pipBody
        $natId=Resource-Id 'Microsoft.Network/natGateways' "nat-ampls-$side"
        $natBody=New-Body @{publicIpAddresses=@(@{id=$pipId});idleTimeoutInMinutes=4};$natBody.sku=@{name='Standard'}
        $null=Ensure-Resource $natId $script:Api.Network $natBody
        $nsgId=Resource-Id 'Microsoft.Network/networkSecurityGroups' "nsg-ampls-$side"
        $null=Ensure-Resource $nsgId $script:Api.Network (New-Body @{securityRules=@(@{name='DenyInboundManagement';properties=@{priority=100;direction='Inbound';access='Deny';protocol='Tcp';sourcePortRange='*';destinationPortRanges=@('22','3389');sourceAddressPrefix='*';destinationAddressPrefix='*'}})})
        $vnetId=Resource-Id 'Microsoft.Network/virtualNetworks' "vnet-ampls-$side"
        $subnets=@(@{name='workload';properties=@{addressPrefix="$prefix.1.0/24";defaultOutboundAccess=$false;natGateway=@{id=$natId};networkSecurityGroup=@{id=$nsgId}}})
        if ($side -eq 'inside') { $subnets+=@{name='private-endpoints';properties=@{addressPrefix="$prefix.2.0/24";privateEndpointNetworkPolicies='Disabled';defaultOutboundAccess=$false}} }
        $null=Ensure-Resource $vnetId $script:Api.Network (New-Body @{addressSpace=@{addressPrefixes=@("$prefix.0.0/16")};subnets=$subnets})
        $nicId=Resource-Id 'Microsoft.Network/networkInterfaces' "nic-ampls-$side"
        $null=Ensure-Resource $nicId $script:Api.Network (New-Body @{ipConfigurations=@(@{name='primary';properties=@{privateIPAllocationMethod='Dynamic';subnet=@{id="$vnetId/subnets/workload"}}})})
        $vmId=Resource-Id 'Microsoft.Compute/virtualMachines' "vm-ampls-$side"
        $vmBody=New-Body @{hardwareProfile=@{vmSize=$script:VmSize};storageProfile=@{imageReference=@{publisher='Canonical';offer='ubuntu-24_04-lts';sku='server';version='latest'};osDisk=@{name="os-ampls-$side";createOption='FromImage';deleteOption='Delete';managedDisk=@{storageAccountType='Standard_LRS'}}};osProfile=@{computerName="vm-ampls-$side";adminUsername='labadmin';linuxConfiguration=@{disablePasswordAuthentication=$true;provisionVMAgent=$true;ssh=@{publicKeys=@(@{path='/home/labadmin/.ssh/authorized_keys';keyData=$publicKey})}}};networkProfile=@{networkInterfaces=@(@{id=$nicId;properties=@{deleteOption='Delete'}})}}
        $vmBody.identity=@{type='SystemAssigned, UserAssigned';userAssignedIdentities=@{$uamiId=@{}}}
        $null=Ensure-Resource $vmId $script:Api.Compute $vmBody
    }
    $lawId=Resource-Id 'Microsoft.OperationalInsights/workspaces' 'law-ampls-deepdive-we'
    $law=Ensure-Resource $lawId $script:Api.Logs (New-Body @{sku=@{name='PerGB2018'};retentionInDays=30;publicNetworkAccessForIngestion='Enabled';publicNetworkAccessForQuery='Enabled'})
    $dceId=Resource-Id 'Microsoft.Insights/dataCollectionEndpoints' 'dce-ampls-we'
    $dce=Ensure-Resource $dceId $script:Api.Insights (New-Body @{networkAcls=@{publicNetworkAccess='Enabled'}})
    $columns=@(@{name='TimeGenerated';type='datetime'},@{name='RunKey';type='string'},@{name='Stage';type='string'},@{name='Source';type='string'},@{name='Marker';type='string'})
    $null=Ensure-Resource "$lawId/tables/PrivateLinkProof_CL" '2022-10-01' @{properties=@{schema=@{name='PrivateLinkProof_CL';columns=$columns};retentionInDays=30;totalRetentionInDays=30}} -Child
    $dcrId=Resource-Id 'Microsoft.Insights/dataCollectionRules' 'dcr-ampls-probe'
    $body=New-Body @{dataCollectionEndpointId=$dceId;streamDeclarations=@{'Custom-PrivateLinkProof_CL'=@{columns=$columns}};destinations=@{logAnalytics=@(@{name='law';workspaceResourceId=$lawId})};dataFlows=@(@{streams=@('Custom-PrivateLinkProof_CL');destinations=@('law');transformKql='source';outputStream='Custom-PrivateLinkProof_CL'})}
    $dcr=Ensure-Resource $dcrId $script:Api.Dcr $body
    foreach ($rolePair in @(@{scope=$lawId;role='Log Analytics Reader'},@{scope=$dcrId;role='Monitoring Metrics Publisher'})) {
        $defs=@(Invoke-Az @('role','definition','list','--name',$rolePair.role))
        if ($defs.Count -ne 1) { throw 'Role definition resolution failed.' }
        $keyName='Role-'+$rolePair.role
        if (-not $script:State[$keyName]) { $script:State[$keyName]=[guid]::NewGuid().ToString();Save-State }
        $roleId="$($rolePair.scope)/providers/Microsoft.Authorization/roleAssignments/$($script:State[$keyName])"
        $null=Ensure-Resource $roleId '2022-04-01' @{properties=@{roleDefinitionId=$defs[0].id;principalId=$identity.properties.principalId;principalType='ServicePrincipal'}} -Child
    }
    $script:State.WorkspaceId=$lawId;$script:State.WorkspaceGuid=$law.properties.customerId
    $script:State.DceId=$dceId;$script:State.IngestionEndpoint=$dce.properties.logsIngestion.endpoint
    $script:State.ConfigEndpoint=$dce.properties.configurationAccess.endpoint
    $script:State.DcrId=$dcrId;$script:State.ImmutableId=$dcr.properties.immutableId
    Save-State
}

function Invoke-Guest($Side,$Shell,$Label) {
    $file=Join-Path $script:Runtime "guest-command-$Label-$Side-$([datetime]::UtcNow.Ticks).sh"
    [IO.File]::WriteAllText($file,($Shell -replace "`r`n","`n"),[Text.UTF8Encoding]::new($false))
    for ($attempt=0;$attempt -lt 12;$attempt++) {
        try {
            $raw=Invoke-Az @('vm','run-command','invoke','--resource-group',$script:Rg,'--name',"vm-ampls-$Side",'--command-id','RunShellScript','--scripts',"@$file")
            break
        } catch {
            if ($attempt -eq 11 -or $_ -notmatch 'Run command extension execution is in progress|TooManyRequests|429|ServiceUnavailable') { throw }
            Start-Sleep 15
        }
    }
    Write-PrivateJson ("guest-$Label-$Side.json") $raw
    $message=(@($raw.value | ForEach-Object { $_.message }) -join "`n")
    if ($message -match '(?m)^AMPLS_RESULT=(.+)$') { return ($Matches[1] | ConvertFrom-Json -AsHashtable -Depth 50) }
    throw "Guest returned no result envelope for $Label/$Side. Inspect private guest output."
}
function Invoke-Probe($Side,$Label,$Marker,[switch]$QueryOnly,[string]$Query) {
    if (-not $Query) { $Query="PrivateLinkProof_CL | where RunKey == '$($script:State.Session)' and Marker == '$Marker' | summarize Records=count() by Source, Stage, Marker" }
    $config=@{clientId=$script:State.Identity.properties.clientId;source=$Side;stage=$Label;runKey=$script:State.Session;marker=$Marker;send=(-not $QueryOnly);ingestionUrl="$($script:State.IngestionEndpoint)/dataCollectionRules/$($script:State.ImmutableId)/streams/Custom-PrivateLinkProof_CL?api-version=2023-01-01";queryUrl="https://api.loganalytics.io/v1/workspaces/$($script:State.WorkspaceGuid)/query";query=$Query;dnsHosts=@(([uri]$script:State.IngestionEndpoint).Host,([uri]$script:State.ConfigEndpoint).Host,'api.loganalytics.io',"$($script:State.WorkspaceGuid).ods.opinsights.azure.com","$($script:State.WorkspaceGuid).oms.opinsights.azure.com")}
    $encoded=[Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes(($config | ConvertTo-Json -Depth 30 -Compress)))
    $python=Get-Content (Join-Path $script:Bundle 'guest/probe.py') -Raw
    $shell="set -eu`nsudo resolvectl flush-caches || true`npython3 - '$encoded' <<'PYAMPLS'`n$python`nPYAMPLS`n"
    $result=Invoke-Guest $Side $shell $Label
    Write-PrivateJson "probe-$Label-$Side.json" $result
    return $result
}
function Get-QueryRows($Probe) {
    if ($Probe.query.status -ne 200 -or -not $Probe.query.body['tables']) { return @() }
    $rows=[Collections.Generic.List[object]]::new()
    foreach ($table in $Probe.query.body.tables) {
        foreach ($row in $table.rows) {
            $record=@{}
            for ($i=0;$i -lt $table.columns.Count;$i++) { $record[$table.columns[$i].name]=$row[$i] }
            $rows.Add($record)
        }
    }
    return ,$rows.ToArray()
}
function Assert-PrivatePeer($Probe) {
    foreach ($request in @('ingestion','query')) {
        if ($Probe[$request] -and $Probe[$request].peer -notin $script:State.PrivateIps) { throw "The $request socket did not connect to an AMPLS private endpoint address." }
    }
    foreach ($hostName in $Probe.dns.Keys) {
        if (-not $Probe.dns[$hostName].Count) { throw 'Endpoint DNS did not resolve.' }
        foreach ($ip in $Probe.dns[$hostName]) { if ($ip -notin $script:State.PrivateIps) { throw "Endpoint DNS resolved outside the AMPLS NIC: $hostName" } }
    }
}
function Wait-Marker($Side,$Label,$Marker,[switch]$Private) {
    $deadline=[datetime]::UtcNow.AddMinutes($script:Timeout)
    do {
        $probe=Invoke-Probe $Side $Label $Marker -QueryOnly
        $rows=Get-QueryRows $probe
        if ($probe.query.status -eq 200 -and @($rows | Where-Object { $_.Marker -eq $Marker -and [int]$_.Records -gt 0 }).Count) {
            if ($Private) { Assert-PrivatePeer $probe }
            return $probe
        }
        if ($probe.query.status -eq 401) { throw 'Query authentication failure is not a network proof.' }
        Start-Sleep 15
    } while ([datetime]::UtcNow -lt $deadline)
    throw 'Marker materialization timeout; no pass evidence generated.'
}
function Send-SuccessfulProbe($Side,$Label,$Marker,[switch]$Private) {
    $deadline=[datetime]::UtcNow.AddMinutes($script:Timeout)
    do {
        $probe=Invoke-Probe $Side $Label $Marker
        if ($probe.ingestion.status -eq 204 -and $probe.query.status -eq 200) {
            if ($Private) { Assert-PrivatePeer $probe }
            return $probe
        }
        if ($probe.ingestion.status -eq 401 -or $probe.query.status -eq 401) { throw 'Authentication failure.' }
        Start-Sleep 15
    } while ([datetime]::UtcNow -lt $deadline)
    throw 'Authenticated successful ingestion/query probe timed out.'
}
function Record-Proof($Name,$Details) {
    $script:State.Evidence[$Name]=@{Passed=$true;ObservedUtc=[datetime]::UtcNow.ToString('o');Details=$Details}
    Save-State
}
function New-ProbeMarker($Label) { "$Label-$([Convert]::ToHexString([Security.Cryptography.RandomNumberGenerator]::GetBytes(6)).ToLowerInvariant())" }
function Invoke-PublicBaseline {
    Require-Stage Foundation
    if ($script:State.Locked -or $script:State['PrivateLinkStarted'] -or $script:State['AmplsId'] -or $script:State.Completed['PrivateLink']) {
        if ($script:State.Evidence['PublicBaseline']) { Write-Host 'Retaining observed public-baseline evidence; public access will not be reopened.';return }
        throw 'Fresh deployment required for public baseline.'
    }
    foreach ($side in @('inside','outside')) {
        $marker="baseline-$side-$($script:State.Session)"
        $null=Send-SuccessfulProbe $side 'PublicBaseline' $marker
        $null=Wait-Marker $side 'PublicBaseline' $marker
    }
    Record-Proof PublicBaseline @{InsideIngestion=204;OutsideIngestion=204;InsideQuery=200;OutsideQuery=200;BothMarkersFound=$true;SameProbeIdentity=$true}
}
function Restore-MonitorDnsLink {
    $zoneId=Resource-Id 'Microsoft.Network/privateDnsZones' $script:Zones[0]
    $body=New-Body @{registrationEnabled=$false;virtualNetwork=@{id=(Resource-Id 'Microsoft.Network/virtualNetworks' 'vnet-ampls-inside')}} 'global'
    $null=Invoke-Arm PUT "$zoneId/virtualNetworkLinks/inside" $script:Api.Dns $body
    $null=Wait-Resource "$zoneId/virtualNetworkLinks/inside" $script:Api.Dns
    $script:State.DnsDrillActive=$false;Save-State
}
function Invoke-PrivateLink {
    Require-Stage PublicBaseline
    $script:State.PrivateLinkStarted=$true;Save-State
    $amplsId=Resource-Id 'Microsoft.Insights/privateLinkScopes' 'ampls-deepdive'
    $mode=if ($script:State.Locked) {'PrivateOnly'} else {'Open'}
    $null=Ensure-Resource $amplsId $script:Api.Ampls (New-Body @{accessModeSettings=@{ingestionAccessMode=$mode;queryAccessMode=$mode}} 'global')
    foreach ($entry in @(@{name='workspace';id=$script:State.WorkspaceId},@{name='dce';id=$script:State.DceId})) {
        $null=Ensure-Resource "$amplsId/scopedResources/$($entry.name)" $script:Api.Ampls @{properties=@{kind='Resource';linkedResourceId=$entry.id}} -Child
    }
    $configs=@()
    foreach ($zone in $script:Zones) {
        $zoneId=Resource-Id 'Microsoft.Network/privateDnsZones' $zone
        $null=Ensure-Resource $zoneId $script:Api.Dns (New-Body @{} 'global')
        $null=Ensure-Resource "$zoneId/virtualNetworkLinks/inside" $script:Api.Dns (New-Body @{registrationEnabled=$false;virtualNetwork=@{id=(Resource-Id 'Microsoft.Network/virtualNetworks' 'vnet-ampls-inside')}} 'global')
        $configs+=@{name=($zone -replace '\.','-');properties=@{privateDnsZoneId=$zoneId}}
    }
    $peId=Resource-Id 'Microsoft.Network/privateEndpoints' 'pe-ampls'
    $pe=Ensure-Resource $peId $script:Api.Network (New-Body @{subnet=@{id="$(Resource-Id 'Microsoft.Network/virtualNetworks' 'vnet-ampls-inside')/subnets/private-endpoints"};privateLinkServiceConnections=@(@{name='ampls';properties=@{privateLinkServiceId=$amplsId;groupIds=@('azuremonitor')}})})
    if ($pe.properties.privateLinkServiceConnections[0].properties.privateLinkServiceConnectionState.status -ne 'Approved') { throw 'Private endpoint connection not approved.' }
    $null=Ensure-Resource "$peId/privateDnsZoneGroups/default" $script:Api.Network @{properties=@{privateDnsZoneConfigs=$configs}} -Child
    $nic=Invoke-Arm GET $pe.properties.networkInterfaces[0].id $script:Api.Network
    $script:State.PrivateIps=@($nic.properties.ipConfigurations | ForEach-Object {$_.properties.privateIPAddress})
    $script:State.AmplsId=$amplsId;$script:State.PrivateEndpointId=$peId;Save-State
    # This stage proves endpoint routing. End-to-end ingestion is asserted AFTER
    # all resource-level and scope-level access settings are applied in Lockdown.
    $deadline=[datetime]::UtcNow.AddMinutes($script:Timeout)
    do {
        $probe=Invoke-Probe inside PrivateLink (New-ProbeMarker 'private-link')
        try { Assert-PrivatePeer $probe; $routed=$true } catch { $routed=$false }
        if ($routed) {
            Record-Proof PrivateLink @{Approved=$true;AllEndpointDnsMatchesPrivateNic=$true;PrivateSocketVerified=$true;PrivateDnsZoneCount=5;RoutingOnlyCheck=$true;ObservedIngestionStatus=$probe.ingestion.status;ObservedQueryStatus=$probe.query.status}
            return
        }
        Start-Sleep 15
    } while ([datetime]::UtcNow -lt $deadline)
    throw 'Private endpoint DNS/socket routing did not converge.'
}
function Invoke-Lockdown {
    Require-Stage PrivateLink
    # Mark intent BEFORE mutating. A partial failure must never allow a baseline reset.
    $script:State.Locked=$true;Save-State
    $null=Invoke-Arm PATCH $script:State.WorkspaceId $script:Api.Logs @{properties=@{publicNetworkAccessForIngestion='Disabled';publicNetworkAccessForQuery='Disabled'}}
    $null=Invoke-Arm PUT $script:State.DceId $script:Api.Insights (New-Body @{networkAcls=@{publicNetworkAccess='Disabled'}})
    $null=Invoke-Arm PUT $script:State.AmplsId $script:Api.Ampls (New-Body @{accessModeSettings=@{ingestionAccessMode='PrivateOnly';queryAccessMode='PrivateOnly'}} 'global')
    $null=Wait-Resource $script:State.DceId $script:Api.Insights
    Record-Proof Lockdown @{WorkspacePublicIngestion='Disabled';WorkspacePublicQuery='Disabled';DcePublicAccess='Disabled';AmplsIngestion='PrivateOnly';AmplsQuery='PrivateOnly'}
}
function Invoke-AMA {
    Require-Stage Lockdown
    $vmId=Resource-Id 'Microsoft.Compute/virtualMachines' 'vm-ampls-inside'
    $dcrId=Resource-Id 'Microsoft.Insights/dataCollectionRules' 'dcr-ampls-syslog'
    $body=New-Body @{dataCollectionEndpointId=$script:State.DceId;dataSources=@{syslog=@(@{name='workshopSyslog';streams=@('Microsoft-Syslog');facilityNames=@('local0');logLevels=@('Notice','Warning','Error','Critical','Alert','Emergency')})};destinations=@{logAnalytics=@(@{name='law';workspaceResourceId=$script:State.WorkspaceId})};dataFlows=@(@{streams=@('Microsoft-Syslog');destinations=@('law')})}
    $body.kind='Linux'
    $null=Ensure-Resource $dcrId $script:Api.Dcr $body
    $null=Ensure-Resource "$vmId/providers/Microsoft.Insights/dataCollectionRuleAssociations/configurationAccessEndpoint" $script:Api.Insights @{properties=@{dataCollectionEndpointId=$script:State.DceId}} -Child
    $null=Ensure-Resource "$vmId/providers/Microsoft.Insights/dataCollectionRuleAssociations/workshop-syslog" $script:Api.Insights @{properties=@{dataCollectionRuleId=$dcrId}} -Child
    $extension=New-Body @{publisher='Microsoft.Azure.Monitor';type='AzureMonitorLinuxAgent';typeHandlerVersion='1.0';autoUpgradeMinorVersion=$true;enableAutomaticUpgrade=$true}
    $null=Ensure-Resource "$vmId/extensions/AzureMonitorLinuxAgent" $script:Api.Compute $extension
    $marker=New-ProbeMarker 'ampls-syslog'
    $shell=@"
set -eu
command -v logger >/dev/null
sudo systemctl is-active rsyslog >/dev/null
logger -p local0.notice -t ampls-workshop '$marker'
echo 'AMPLS_RESULT={"emitted":true}'
"@
    $null=Invoke-Guest inside $shell 'syslog-marker'
    $since=[datetime]::UtcNow.AddMinutes(-3).ToString('o')
    $kql="union (Heartbeat | where TimeGenerated > datetime($since) and Computer == 'vm-ampls-inside' | summarize Records=count() | extend Proof='Heartbeat'), (Syslog | where TimeGenerated > datetime($since) and SyslogMessage contains '$marker' | summarize Records=count() | extend Proof='Syslog') | project Proof, Records"
    $deadline=[datetime]::UtcNow.AddMinutes($script:Timeout)
    do {
        $probe=Invoke-Probe inside AMA $marker -QueryOnly -Query $kql
        $rows=Get-QueryRows $probe
        if ($probe.query.status -eq 200 -and @($rows | Where-Object { $_.Proof -eq 'Heartbeat' -and [int]$_.Records -gt 0 }).Count -and @($rows | Where-Object { $_.Proof -eq 'Syslog' -and [int]$_.Records -gt 0 }).Count) {
            Assert-PrivatePeer $probe
            $heartbeatCount=[int](@($rows | Where-Object Proof -eq 'Heartbeat')[0].Records)
            $syslogCount=[int](@($rows | Where-Object Proof -eq 'Syslog')[0].Records)
            Record-Proof AMA @{FreshHeartbeat=$true;HeartbeatRecords=$heartbeatCount;UniqueSyslogMarker=$true;SyslogRecords=$syslogCount;PrivateQuery=$true;DcrAndDceAssociated=$true};return
        }
        $null=Invoke-Guest inside $shell 'syslog-marker'
        Start-Sleep 15
    } while ([datetime]::UtcNow -lt $deadline)
    throw 'AMA Heartbeat or Syslog proof did not materialize.'
}
function Test-NetworkDenial($Response) {
    if ($Response.status -ne 403) { return $false }
    $text=$Response.body | ConvertTo-Json -Depth 20 -Compress
    # The Logs query endpoint can wrap its explicit network denial in an
    # InsufficientAccessError. Accept the specific nested network code/message,
    # never the generic outer authorization error on its own.
    if ($text -match 'NspValidationFailedError' -and $text -match '(?i)access.*denied.*public networks') { return $true }
    return ($text -match '(?i)public.*(access|ingestion|query).*(disabled|denied|not allowed)|network.*(access|isolation)|private.?link|IP.*(not allowed|forbidden)|NetworkAccessDenied')
}
function Get-ServiceErrorCode($Response,[switch]$Inner) {
    $body=$Response.body
    if ($body -is [Collections.IDictionary] -and $body['error'] -is [Collections.IDictionary]) {
        $errorObject=$body.error
        if ($Inner) { $errorObject=$errorObject['innererror'] }
        if ($errorObject -is [Collections.IDictionary]) {
            $code=[string]$errorObject['code']
            if ($code -cmatch '^[A-Za-z][A-Za-z0-9_.-]{0,79}$') { return $code }
        }
    }
    return 'NoStructuredCode'
}
function Invoke-TestBoundary {
    Require-Stage AMA
    $marker=New-ProbeMarker 'boundary-inside'
    $null=Send-SuccessfulProbe inside TestBoundary $marker -Private
    $null=Wait-Marker inside TestBoundary $marker -Private
    $deadline=[datetime]::UtcNow.AddMinutes($script:Timeout)
    do {
        $outside=Invoke-Probe outside TestBoundary (New-ProbeMarker 'boundary-outside')
        if ((Test-NetworkDenial $outside.ingestion) -and (Test-NetworkDenial $outside.query)) {
            if ($outside.ingestion.peer -in $script:State.PrivateIps -or $outside.query.peer -in $script:State.PrivateIps) { throw 'Outside traffic unexpectedly used the private endpoint.' }
            Record-Proof TestBoundary @{InsideIngestion=204;InsideQuery=200;InsideMarkerFound=$true;OutsideIngestion=$outside.ingestion.status;OutsideQuery=$outside.query.status;OutsideIngestionErrorCode=(Get-ServiceErrorCode $outside.ingestion);OutsideQueryErrorCode=(Get-ServiceErrorCode $outside.query);OutsideQueryInnerErrorCode=(Get-ServiceErrorCode $outside.query -Inner);OutsideNetworkDenials=$true;SameIdentityAsBaseline=$true;PrivateSocketVerified=$true};return
        }
        if ($outside.ingestion.status -eq 401 -or $outside.query.status -eq 401) { throw 'Outside authentication failure; cannot accept network proof.' }
        Start-Sleep 15
    } while ([datetime]::UtcNow -lt $deadline)
    throw 'Outside requests did not produce explicit service-side network denials. Inspect private response bodies; do not label generic errors as success.'
}
function Invoke-TestDnsRecovery {
    Require-Stage TestBoundary
    $script:State.DnsDrillActive=$true;Save-State
    $zoneId=Resource-Id 'Microsoft.Network/privateDnsZones' $script:Zones[0]
    $broken=$false
    try {
        $null=Invoke-Arm DELETE "$zoneId/virtualNetworkLinks/inside" $script:Api.Dns
        $deadline=[datetime]::UtcNow.AddMinutes($script:Timeout)
        do {
            $probe=Invoke-Probe inside DnsBroken (New-ProbeMarker 'dns-broken')
            if ((Test-NetworkDenial $probe.ingestion) -and (Test-NetworkDenial $probe.query) -and $probe.ingestion.peer -notin $script:State.PrivateIps -and $probe.query.peer -notin $script:State.PrivateIps) { $broken=$true;break }
            Start-Sleep 15
        } while ([datetime]::UtcNow -lt $deadline)
    } finally { Restore-MonitorDnsLink }
    if (-not $broken) { throw 'DNS break did not produce the required denial; link restored.' }
    $marker=New-ProbeMarker 'dns-recovered'
    $null=Send-SuccessfulProbe inside DnsRecovered $marker -Private
    $null=Wait-Marker inside DnsRecovered $marker -Private
    Record-Proof TestDnsRecovery @{BrokenDnsNetworkDenial=$true;DnsLinkRestored=$true;RecoveredIngestion=204;RecoveredQuery=200;RecoveredMarkerFound=$true}
}
function Invoke-Validate {
    foreach ($name in @('PublicBaseline','PrivateLink','Lockdown','AMA','TestBoundary','TestDnsRecovery')) { Require-Stage $name; if (-not $script:State.Evidence[$name].Passed) { throw "Missing observed evidence: $name" } }
    if ($script:State['DnsDrillActive']) { Restore-MonitorDnsLink;throw 'Recovered interrupted DNS drill; rerun TestDnsRecovery.' }
    $law=Invoke-Arm GET $script:State.WorkspaceId $script:Api.Logs
    $dce=Invoke-Arm GET $script:State.DceId $script:Api.Insights
    $ampls=Invoke-Arm GET $script:State.AmplsId $script:Api.Ampls
    if ($law.properties.publicNetworkAccessForIngestion -ne 'Disabled' -or $law.properties.publicNetworkAccessForQuery -ne 'Disabled' -or $dce.properties.networkAcls.publicNetworkAccess -ne 'Disabled' -or $ampls.properties.accessModeSettings.ingestionAccessMode -ne 'PrivateOnly' -or $ampls.properties.accessModeSettings.queryAccessMode -ne 'PrivateOnly') { throw 'Final public-access drift.' }
    foreach ($side in @('inside','outside')) {
        $nic=Invoke-Arm GET (Resource-Id 'Microsoft.Network/networkInterfaces' "nic-ampls-$side") $script:Api.Network
        if (@($nic.properties.ipConfigurations | Where-Object {$_.properties['publicIPAddress']}).Count) { throw 'Unexpected VM public IP.' }
        $vnet=Invoke-Arm GET (Resource-Id 'Microsoft.Network/virtualNetworks' "vnet-ampls-$side") $script:Api.Network
        if (@($vnet.properties.virtualNetworkPeerings).Count) { throw 'Unexpected VNet peering.' }
        $workload=@($vnet.properties.subnets | Where-Object name -eq 'workload')[0]
        if ($workload.properties.defaultOutboundAccess -ne $false -or -not $workload.properties.natGateway.id) { throw 'Unexpected outbound configuration.' }
    }
    foreach ($zone in $script:Zones) {
        $links=Invoke-Arm GET "$(Resource-Id 'Microsoft.Network/privateDnsZones' $zone)/virtualNetworkLinks" $script:Api.Dns
        if ($links.value.Count -ne 1 -or $links.value[0].properties.virtualNetwork.id -ne (Resource-Id 'Microsoft.Network/virtualNetworks' 'vnet-ampls-inside')) { throw 'Private DNS link drift.' }
    }
    $insideVmId=Resource-Id 'Microsoft.Compute/virtualMachines' 'vm-ampls-inside'
    $configAssociation=Invoke-Arm GET "$insideVmId/providers/Microsoft.Insights/dataCollectionRuleAssociations/configurationAccessEndpoint" $script:Api.Insights
    $ruleAssociation=Invoke-Arm GET "$insideVmId/providers/Microsoft.Insights/dataCollectionRuleAssociations/workshop-syslog" $script:Api.Insights
    if ($configAssociation.properties.dataCollectionEndpointId -ne $script:State.DceId -or $ruleAssociation.properties.dataCollectionRuleId -ne (Resource-Id 'Microsoft.Insights/dataCollectionRules' 'dcr-ampls-syslog')) { throw 'AMA association drift.' }
    # Recheck delivery now instead of accepting a stale earlier AMA-stage result.
    Invoke-AMA
    $finalMarker=New-ProbeMarker 'final'
    $null=Send-SuccessfulProbe inside Validate $finalMarker -Private
    $null=Wait-Marker inside Validate $finalMarker -Private
    Record-Proof Validate @{FinalPrivateConfiguration=$true;NoVmPublicIps=$true;NoPeerings=$true;FiveInsideOnlyDnsLinks=$true;ExplicitOutbound=$true;AmaAssociationsVerified=$true;FreshHeartbeatAndSyslogRechecked=$true}
}
