From 8a83b15e5feb829549e72d523734f0e657ac831b Mon Sep 17 00:00:00 2001
From: Abigailliang <nliang@microsoft.com>
Date: Thu, 12 Nov 2026 13:51:26 -0800
Subject: [PATCH] fix: isolate additional Linux SSH lifecycle safeguards

Move the previously implemented activation restoration, public-key authentication save/restore and validation, explicit restoration rendering, and shutdown verification out of PR #10021 into its dependent follow-up. Preserve PR #10063 nodePrep placement and current dev changes. Same-node restoration remains a separately reviewed safeguard, not a requirement for the normal SSH update followed by reimage.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: c815829b-19d2-4af1-8269-5d704d14325c
---
 aks-node-controller/parser/parser.go          |   3 +
 aks-node-controller/parser/parser_test.go     |  26 +
 e2e/README.md                                 |  10 +-
 parts/linux/cloud-init/artifacts/README.MD    |  17 +-
 parts/linux/cloud-init/artifacts/cse_cmd.sh   |   1 +
 .../linux/cloud-init/artifacts/cse_config.sh  | 285 +++++++-
 parts/linux/cloud-init/artifacts/cse_main.sh  |   7 +-
 pkg/agent/baker.go                            |   3 +
 pkg/agent/baker_test.go                       |  14 +
 .../cloud-init/artifacts/cse_config_spec.sh   | 684 +++++++++++++++++-
 .../artifacts/cse_main_pis_reconcile_spec.sh  |   6 +-
 .../artifacts/cse_ssh_policy_spec.sh          | 268 +++----
 spec/support/ssh_policy.sh                    |  98 +++
 13 files changed, 1230 insertions(+), 192 deletions(-)
 create mode 100644 spec/support/ssh_policy.sh

diff --git a/aks-node-controller/parser/parser.go b/aks-node-controller/parser/parser.go
index 4015dbde80e..8d10e96a611 100644
--- a/aks-node-controller/parser/parser.go
+++ b/aks-node-controller/parser/parser.go
@@ -47,6 +47,8 @@ func getCSEEnv(ctx context.Context, config *aksnodeconfigv1.Configuration, gpuCo
 	containerdVersion, _ := detectContainerdVersion(ctx)
 	cloudProviderSettings := getCloudProviderSettings(config)
 	isMIGNode := getIsMIGNode(config.GetGpuConfig().GetGpuInstanceProfile(), config.GetGpuConfig().GetMigProfileLayout())
+	// Explicit SSH disablement wins; otherwise allow explicit SSH or Entra-only access.
+	shouldRestoreSSH := !getDisableSSH(config) && (config.GetEnableSsh() || config.GetDisablePubkeyAuth())
 	env := map[string]string{
 		"PROVISION_OUTPUT":                                     "/var/log/azure/cluster-provision-cse-output.log",
 		"CUSTOM_DATA_ONLY_PROVISIONING":                        fmt.Sprintf("%t", config.GetCustomDataOnlyProvisioning()),
@@ -112,6 +114,7 @@ func getCSEEnv(ctx context.Context, config *aksnodeconfigv1.Configuration, gpuCo
 		"RUNC_PACKAGE_URL":                                     config.GetRuncConfig().GetRuncPackageUrl(),
 		"ENABLE_HOSTS_CONFIG_AGENT":                            fmt.Sprintf("%v", config.GetEnableHostsConfigAgent()),
 		"DISABLE_SSH":                                          fmt.Sprintf("%v", getDisableSSH(config)),
+		"SHOULD_RESTORE_SSH":                                   fmt.Sprintf("%v", shouldRestoreSSH),
 		"SHOULD_CONFIGURE_HTTP_PROXY":                          fmt.Sprintf("%v", getShouldConfigureHTTPProxy(config.GetHttpProxyConfig())),
 		"SHOULD_CONFIGURE_HTTP_PROXY_CA":                       fmt.Sprintf("%v", getShouldConfigureHTTPProxyCA(config.GetHttpProxyConfig())),
 		"HTTP_PROXY_TRUSTED_CA":                                removeNewlines(config.GetHttpProxyConfig().GetProxyTrustedCa()),
diff --git a/aks-node-controller/parser/parser_test.go b/aks-node-controller/parser/parser_test.go
index cc2d592d22c..7b1bb30f78b 100644
--- a/aks-node-controller/parser/parser_test.go
+++ b/aks-node-controller/parser/parser_test.go
@@ -21,6 +21,32 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
+func TestBuildCSECmdSSHPolicy(t *testing.T) {
+	for _, test := range []struct {
+		name      string
+		enable    *bool
+		pubkeyOff bool
+		disable   string
+		restore   string
+	}{
+		{name: "unspecified", disable: "false", restore: "false"},
+		{name: "disabled", enable: proto.Bool(false), disable: "true", restore: "false"},
+		{name: "explicitly enabled", enable: proto.Bool(true), disable: "false", restore: "true"},
+		{name: "explicit Entra authentication", pubkeyOff: true, disable: "false", restore: "true"},
+		{name: "disabled takes precedence", enable: proto.Bool(false), pubkeyOff: true, disable: "true", restore: "false"},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			cmd, err := BuildCSECmd(context.Background(), &aksnodeconfigv1.Configuration{
+				EnableSsh: test.enable, DisablePubkeyAuth: proto.Bool(test.pubkeyOff),
+			}, nil)
+			require.NoError(t, err)
+			env := environToMap(cmd.Env)
+			require.Equal(t, test.disable, env["DISABLE_SSH"])
+			require.Equal(t, test.restore, env["SHOULD_RESTORE_SSH"])
+		})
+	}
+}
+
 func TestBuildCSECmdContainerdImagePullSettings(t *testing.T) {
 	for _, version := range []string{"1.7.22", "2.2.4"} {
 		t.Run(version, func(t *testing.T) {
diff --git a/e2e/README.md b/e2e/README.md
index 11bf8a36d4b..4bb88cbfe33 100644
--- a/e2e/README.md
+++ b/e2e/README.md
@@ -30,12 +30,14 @@ and require an updated CSE bundle or VHD to receive this fix.
 
 Linux SSH-disabled scenarios use RunCommand to verify the persistent SSH block,
 restart the same VM through ARM, confirm its boot ID changed, and check the block
-again after kubelet becomes active. CSE stops SSH listeners and uses an AKS-owned
-systemd condition without changing distribution enablement or administrator masks.
+again after kubelet becomes active. CSE records activation, stops and disables SSH
+listeners, and uses an AKS-owned systemd condition without removing administrator masks.
 SSH policy runs only in `nodePrep`, so image preparation leaves SSH defaults
 unchanged. This delays SSH policy enforcement until `nodePrep` on ordinary nodes
-as well. There is no activation history or automatic re-enable path; a code
-rollback alone does not remove a persisted block.
+as well. The follow-up retains activation and public-key authentication snapshots
+for explicit restoration. These are separate safeguards, not a requirement of the
+normal node-pool SSH configuration change, which takes effect through reimage.
+A code rollback alone does not remove a persisted block.
 
 ## Writing and extending scenarios
 
diff --git a/parts/linux/cloud-init/artifacts/README.MD b/parts/linux/cloud-init/artifacts/README.MD
index ff5c67702a2..14ec7150cfe 100644
--- a/parts/linux/cloud-init/artifacts/README.MD
+++ b/parts/linux/cloud-init/artifacts/README.MD
@@ -7,10 +7,10 @@ For CA refresh logging and Linux opt-in certificate diagnostics, see
 
 When `DISABLE_SSH=true` and `nodePrep` runs, `disableSSH` emits structured
 guest-agent events. A persistent systemd condition blocks future starts; existing
-listeners are stopped without changing unit enablement. Stop retry parameters and
-missing-unit handling are unchanged. Policy-installation and stop failures surface
-as CSE exit code 172 after every SSH unit has been attempted. A successful stop
-does not have to clear a unit's existing `failed` state.
+listeners are stopped and disabled after their activation policy is recorded.
+Stop retry parameters and missing-unit handling are unchanged. Policy-installation,
+inspection, stop, disable, or inactive-state verification failures surface as CSE
+exit code 172 after every SSH unit has been attempted.
 Telemetry failures remain warnings, not new
 provisioning failures. Other `systemctlDisableAndStop` callers and
 `logs_to_events` keep their existing telemetry behavior.
@@ -18,7 +18,7 @@ provisioning failures. Other `systemctlDisableAndStop` callers and
 | TaskName | Scope |
 | --- | --- |
 | `AKS.CSE.disableSSH` | One `Started` and one `Completed` event per SSH-disable attempt |
-| `AKS.CSE.disableSSH.unit` | `Started` and `Completed` for each unit reached: `ssh.socket`, `sshd.socket`, `ssh`, `sshd` |
+| `AKS.CSE.disableSSH.unit` | `Started` and `Completed` for each unit handled by `systemctlDisableAndStop`: `ssh.socket`, `sshd.socket`, `ssh`, `sshd` |
 
 The event envelope's `Message` is a JSON string, available as `Context1` in
 `GuestAgentGenericLogs`. Parse it with `parse_json(Context1)`.
@@ -31,13 +31,14 @@ The event envelope's `Message` is a JSON string, available as `Context1` in
 | `ExitCode` | Null for starts; 0/1 for completed units; 0/172 for the completed attempt |
 | `CatExitCode`, `Skipped` | Whether the unit was readable; nonzero `systemctl cat` preserves the existing skip behavior |
 | `StopExitCode` | Stop retry-helper result; null when skipped, not the raw systemctl exit code |
-| `DisableExitCode` | Always null for this stop-only SSH path; retained for schema compatibility |
+| `DisableExitCode` | Disable retry-helper result; null when skipped |
 
 The last three rows apply to completed unit events. The envelope's `Timestamp` and
 `OperationId` contain start and emission times. The schema is kept compact to fit
 the custom-data size limit. A skipped unit can be absent or unreadable;
-it is not evidence of a successful stop. Readable masked units are stopped and
-emit per-unit events without removing their masks. Command exit 0 alone does not prove port
+it is not evidence of a successful stop. Externally masked units are stopped without
+removing their masks and emit only the attempt-level events. Other readable units
+emit per-unit events. Command exit 0 alone does not prove port
 22 is closed: inspect the unit states and use the reboot E2E validation for
 listener checks.
 
diff --git a/parts/linux/cloud-init/artifacts/cse_cmd.sh b/parts/linux/cloud-init/artifacts/cse_cmd.sh
index a08d3b35d6d..ac0f86249c6 100644
--- a/parts/linux/cloud-init/artifacts/cse_cmd.sh
+++ b/parts/linux/cloud-init/artifacts/cse_cmd.sh
@@ -94,6 +94,7 @@ RUNC_PACKAGE_URL={{GetParameter "runcPackageURL"}}
 ENABLE_HOSTS_CONFIG_AGENT="{{EnableHostsConfigAgent}}"
 DISABLE_SSH="{{ShouldDisableSSH}}"
 DISABLE_PUBKEY_AUTH="{{ShouldTurnOffPubkeyAuthSSH}}"
+SHOULD_RESTORE_SSH="{{ShouldRestoreSSH}}"
 SHOULD_CONFIGURE_HTTP_PROXY="{{ShouldConfigureHTTPProxy}}"
 SHOULD_CONFIGURE_HTTP_PROXY_CA="{{ShouldConfigureHTTPProxyCA}}"
 HTTP_PROXY_TRUSTED_CA="{{GetHTTPProxyCA}}"
diff --git a/parts/linux/cloud-init/artifacts/cse_config.sh b/parts/linux/cloud-init/artifacts/cse_config.sh
index 2023f28f47e..8ed9acfc3ef 100755
--- a/parts/linux/cloud-init/artifacts/cse_config.sh
+++ b/parts/linux/cloud-init/artifacts/cse_config.sh
@@ -813,27 +813,197 @@ createKubeManifestDir() {
     mkdir -p $KUBEMANIFESTDIR
 }
 
+# Prefer socket activation when the distribution enables it; aliases resolve to Id.
+getSSHActivation() {
+    local unit load_state enabled_state active_state id candidate=""
+    for unit in ssh.socket sshd.socket ssh.service sshd.service; do
+        load_state=$(systemctl show "$unit" --property=LoadState --value) || return 1
+        case "$load_state" in
+            not-found|masked) continue ;;
+            loaded) ;;
+            *) echo "Cannot inspect SSH unit $unit: $load_state" >&2; return 1 ;;
+        esac
+        enabled_state=$(systemctl show "$unit" --property=UnitFileState --value) || return 1
+        if [ "$enabled_state" = enabled ]; then
+            id=$(systemctl show "$unit" --property=Id --value) || return 1
+            printf '%s enabled\n' "$id"
+            return 0
+        fi
+    done
+    for unit in ssh.socket sshd.socket ssh.service sshd.service; do
+        active_state=$(systemctl show "$unit" --property=ActiveState --value) || return 1
+        if [ "$active_state" = active ]; then
+            id=$(systemctl show "$unit" --property=Id --value) || return 1
+            printf '%s disabled\n' "$id"
+            return 0
+        fi
+    done
+    [ "${1:-false}" = true ] || return 0
+    # Legacy stop+disable leaves no enabled or active endpoint. Capture only
+    # an unambiguous unmasked unit, without inventing enabled boot links.
+    for unit in ssh.socket sshd.socket ssh.service sshd.service; do
+        load_state=$(systemctl show "$unit" --property=LoadState --value) || return 1
+        [ "$load_state" = loaded ] || continue
+        id=$(systemctl show "$unit" --property=Id --value) || return 1
+        if [ -n "$candidate" ] && [ "$candidate" != "$id" ]; then
+            return 0
+        fi
+        candidate="$id"
+    done
+    [ -z "$candidate" ] || printf '%s disabled\n' "$candidate"
+}
+
 installSSHDisablePolicy() {
     local policy_dir="${SSH_POLICY_DIR:-/etc/aks/ssh-policy}"
     local unit_dir="${SSH_SYSTEMD_UNIT_DIR:-/etc/systemd/system}"
-    local unit dropin content
-    mkdir -p "$policy_dir" && chmod 0700 "$policy_dir" || return 1
+    local unit dropin content result=0
+    mkdir -p "$policy_dir" || return 1
+    chmod 0700 "$policy_dir" || return 1
+    if [ ! -f "$policy_dir/activation" ]; then
+        getSSHActivation true > "$policy_dir/activation.tmp" &&
+            mv "$policy_dir/activation.tmp" "$policy_dir/activation" || return 1
+    fi
     touch "$policy_dir/disabled" || return 1
-    # Package scripts can undo masks. A separate condition also blocks dependency starts.
+    # Unlike disable (or a mask undone by openssh package migration), this
+    # condition also blocks dependency and socket starts after a reboot.
     content="[Unit]
 ConditionPathExists=!$policy_dir/disabled"
     for unit in ssh.socket sshd.socket ssh.service sshd.service; do
         dropin="$unit_dir/$unit.d/90-aks-ssh-access.conf"
         if [ -e "$dropin" ] && [ "$(cat "$dropin")" != "$content" ]; then
             echo "Refusing to overwrite an existing SSH policy: $dropin" >&2
-            return 1
+            result=1
+            continue
         fi
         mkdir -p "$unit_dir/$unit.d" &&
             printf '%s\n' "$content" > "$dropin.tmp" &&
             chmod 0644 "$dropin.tmp" &&
-            mv "$dropin.tmp" "$dropin" || return 1
+            mv "$dropin.tmp" "$dropin" || result=1
     done
-    systemctl daemon-reload
+    systemctl daemon-reload || result=1
+    return "$result"
+}
+
+validateSSHActivation() {
+    local unit="$1" load_state accept triggers service
+    local -a services=()
+    load_state=$(systemctl show "$unit" --property=LoadState --value) || return 1
+    if [ "$load_state" != loaded ]; then
+        echo "Cannot restore SSH: $unit is $load_state; external masks are not removed" >&2
+        return 1
+    fi
+    case "$unit" in
+        *.service) return 0 ;;
+        *.socket)
+            accept=$(systemctl show "$unit" --property=Accept --value) || return 1
+            case "$accept" in
+                yes)
+                    # Inspect a template instance without starting a connection handler.
+                    services=("${unit%.socket}@internal.service")
+                    ;;
+                no)
+                    triggers=$(systemctl show "$unit" --property=Triggers --value) || return 1
+                    read -r -a services <<< "$triggers"
+                    ;;
+                *) echo "Cannot restore SSH: invalid Accept state for $unit: $accept" >&2; return 1 ;;
+            esac
+            ;;
+        *) echo "Cannot restore SSH: unsupported activation unit $unit" >&2; return 1 ;;
+    esac
+    if [ "${#services[@]}" -eq 0 ]; then
+        echo "Cannot restore SSH: no backing service reported for $unit" >&2
+        return 1
+    fi
+    for service in "${services[@]}"; do
+        case "$service" in
+            *.service) validateSSHActivation "$service" || return 1 ;;
+            *) echo "Cannot restore SSH: invalid backing service for $unit: $service" >&2; return 1 ;;
+        esac
+    done
+}
+
+restoreSSH() {
+    local policy_dir="${SSH_POLICY_DIR:-/etc/aks/ssh-policy}"
+    local unit activation load_state id candidate="" owned_policy=false legacy_activation=false
+    [ ! -f "$policy_dir/disabled" ] || owned_policy=true
+    if [ "$owned_policy" = true ] && [ ! -f "$policy_dir/activation" ]; then
+        echo "Cannot restore SSH: the owned activation record is missing" >&2
+        return 1
+    fi
+    # Package upgrades may have migrated service activation to a socket while
+    # our condition kept it blocked. Respect that mode before the saved mode.
+    activation=$(getSSHActivation) || return 1
+    if [ -z "$activation" ] && [ "$owned_policy" = true ]; then
+        activation=$(cat "$policy_dir/activation") || return 1
+    fi
+    # Legacy CSE discarded the original service/socket mode. An explicit enable
+    # may select a unique unmasked service, but never invent the previous mode.
+    if [ -z "$activation" ]; then
+        for unit in ssh.service sshd.service; do
+            load_state=$(systemctl show "$unit" --property=LoadState --value) || return 1
+            case "$load_state" in
+                not-found) continue ;;
+                loaded) ;;
+                *) echo "Cannot restore SSH: $unit is $load_state; external masks are not removed" >&2; return 1 ;;
+            esac
+            id=$(systemctl show "$unit" --property=Id --value) || return 1
+            if [ -n "$candidate" ] && [ "$candidate" != "$id" ]; then
+                echo "Cannot restore SSH: multiple distinct services are installed" >&2
+                return 1
+            fi
+            candidate="$id"
+        done
+        if [ -n "$candidate" ]; then
+            activation="$candidate enabled"
+            legacy_activation=true
+        fi
+    fi
+    read -r unit activation <<< "$activation"
+    case "$unit:$activation" in
+        ssh.service:enabled|sshd.service:enabled|ssh.socket:enabled|sshd.socket:enabled|\
+        ssh.service:disabled|sshd.service:disabled|ssh.socket:disabled|sshd.socket:disabled) ;;
+        *) echo "Cannot restore SSH: no valid saved activation mode" >&2; return 1 ;;
+    esac
+    validateSSHActivation "$unit" || return 1
+    if [ "$owned_policy" = true ] || [ "$legacy_activation" = true ]; then
+        [ "$activation" != enabled ] || systemctl enable "$unit" || return 1
+    fi
+    if [ "$owned_policy" = true ]; then
+        rm -f "$policy_dir/disabled" || return 1
+    fi
+    # Keep the owned drop-ins: without the marker their conditions allow starts.
+    if ! systemctl start "$unit" ||
+        ! systemctl is-active --quiet "$unit"; then
+        if [ "$owned_policy" = true ]; then
+            touch "$policy_dir/disabled" || return 1
+            systemctl stop "$unit" || return 1
+        fi
+        echo "Failed to restore SSH activation through $unit" >&2
+        return 1
+    fi
+    [ "$owned_policy" != true ] || rm -f "$policy_dir/activation"
+}
+
+applySSHPolicy() {
+    local policy_dir="${SSH_POLICY_DIR:-/etc/aks/ssh-policy}"
+    if [ "${DISABLE_SSH:-}" = true ]; then
+        disableSSH
+    elif [ "${DISABLE_PUBKEY_AUTH:-}" = true ]; then
+        logs_to_events "AKS.CSE.disableSSHPubkeyAuth" disableSSHPubkeyAuth
+    elif [ "${SHOULD_RESTORE_SSH:-}" = true ]; then
+        if [ -f "$policy_dir/pubkey-auth" ]; then
+            local pubkey_auth
+            pubkey_auth=$(cat "$policy_dir/pubkey-auth") || return "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+            configureSSHPubkeyAuth "$pubkey_auth" || return "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+            if [ -f "$policy_dir/disabled" ]; then
+                restoreSSH || return "$ERR_DISABLE_SSH"
+            fi
+            reloadSSHAuthentication false || return "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+            rm -f "$policy_dir/pubkey-auth" || return "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+        elif [ -f "$policy_dir/disabled" ]; then
+            restoreSSH || return "$ERR_DISABLE_SSH"
+        fi
+    fi
 }
 
 disableSSH() {
@@ -845,20 +1015,49 @@ disableSSH() {
     installSSHDisablePolicy || { echo "Failed to install persistent SSH disable policy" >&2; result=$ERR_DISABLE_SSH; }
     # Ubuntu uses ssh, Azure Linux uses sshd, and ACL also uses sshd.socket.
     # Close the listening socket first so no new connection is accepted mid-teardown.
-    # Keep going after a failure so a later unit still gets stopped; on ACL the
+    # Keep going after a failure so a later unit still gets disabled, on ACL the
     # socket owns port 22 and skipping it would leave SSH reachable.
+    local load_state active_state
     for unit in ssh.socket sshd.socket ssh sshd; do
-        systemctlStop "$unit" "$attempt_id" || result=$ERR_DISABLE_SSH
+        if ! load_state=$(systemctl show "$unit" --property=LoadState --value); then
+            echo "Failed to inspect SSH unit $unit" >&2
+            result=$ERR_DISABLE_SSH
+            continue
+        fi
+        case "$load_state" in
+            loaded|not-found)
+                systemctlDisableAndStop "$unit" "$attempt_id" || result=$ERR_DISABLE_SSH
+                ;;
+            masked)
+                # A masked unit can still be running. Never remove external masks.
+                systemctl_stop 20 5 25 "$unit" || result=$ERR_DISABLE_SSH
+                ;;
+            *)
+                echo "Cannot disable SSH unit $unit: $load_state" >&2
+                result=$ERR_DISABLE_SSH
+                ;;
+        esac
+        [ "$load_state" != not-found ] || continue
+        active_state=$(systemctl show "$unit" --property=ActiveState --value) || result=$ERR_DISABLE_SSH
+        if [ "$active_state" != inactive ]; then
+            echo "SSH unit $unit is not inactive: $active_state" >&2
+            result=$ERR_DISABLE_SSH
+        fi
     done
     emitSSHDisableEvent "$attempt_id" Completed "" "$start_time" "$result" '{}' ||
         echo "WARNING: SSH disable telemetry failed" >&2
     [ "$result" -eq 0 ] || exit "$result"
 }
 
-disableSSHPubkeyAuth() {
+configureSSHPubkeyAuth() {
+  local desired="$1"
+  case "$desired" in
+    yes|no) ;;
+    *) echo "Invalid SSH public-key authentication policy: $desired" >&2; return 1 ;;
+  esac
   local SSHD_CONFIG="${SSHD_CONFIG_FILE:-/etc/ssh/sshd_config}"
   local TMP
-  TMP="$(mktemp)"
+  TMP="$(mktemp)" || return 1
 
   # AAD SSH extension will append following section to the end of sshd_config,
   # so we need to check the "Match" section, and only update "PubkeyAuthentication" outside of it.
@@ -867,7 +1066,7 @@ disableSSHPubkeyAuth() {
   # PubkeyAuthentication yes
   # AuthorizedKeysCommand /usr/sbin/aad_certhandler %u %k
   # AuthorizedKeysCommandUser root
-  awk -v desired="no" '
+  awk -v desired="$desired" '
     BEGIN { in_match=0; replaced=0; inserted=0 }
     /^Match([[:space:]]|$)/ {
       if (!replaced && !inserted) { print "PubkeyAuthentication " desired; inserted=1 }
@@ -878,24 +1077,64 @@ disableSSHPubkeyAuth() {
     }
     { print }
     END { if (!replaced && !inserted) print "PubkeyAuthentication " desired }
-  ' "$SSHD_CONFIG" > "$TMP"
-
-  local ssh_service="sshd.service"
-  if systemctl cat ssh.service >/dev/null 2>&1; then
-    ssh_service="ssh.service"
+  ' "$SSHD_CONFIG" > "$TMP" || { rm -f "$TMP"; return 1; }
+
+  # Ubuntu removes this runtime directory when ssh.service stops. Validation
+  # must not start the server merely to recreate its privilege-separation path.
+  install -d -m 0755 -o root -g root /run/sshd || { rm -f "$TMP"; return 1; }
+
+  # Validate syntax and the effective global policy, including Include files.
+  sshd -t -f "$TMP" || { rm -f "$TMP"; return 1; }
+  local effective
+  effective=$(sshd -T -f "$TMP") || { rm -f "$TMP"; return 1; }
+  if [ "$(awk 'tolower($1) == "pubkeyauthentication" { print $2 }' <<< "$effective")" != "$desired" ]; then
+    echo "SSH configuration does not enforce the requested public-key policy" >&2
+    rm -f "$TMP"
+    return 1
   fi
-  systemctl is-active --quiet "$ssh_service" || systemctl start "$ssh_service" || exit $ERR_CONFIG_PUBKEY_AUTH_SSH
-
-  # Validate the candidate config
-  sshd -t -f "$TMP" || { rm -f "$TMP"; exit $ERR_CONFIG_PUBKEY_AUTH_SSH; }
 
   # Replace the original with the candidate (permissions 600, owned by root)
   # Mode 0600 is required by CIS Benchmark control 5.1.1
-  install -m 0600 -o root -g root "$TMP" "$SSHD_CONFIG"
+  install -m 0600 -o root -g root "$TMP" "$SSHD_CONFIG" || { rm -f "$TMP"; return 1; }
   rm -f "$TMP"
+}
+
+reloadSSHAuthentication() {
+  local start_if_inactive="${1:-true}"
+  local ssh_service=sshd.service
+  if systemctl cat ssh.service >/dev/null 2>&1; then
+    ssh_service=ssh.service
+  fi
+  if ! systemctl is-active --quiet "$ssh_service"; then
+    [ "$start_if_inactive" != false ] || return 0
+    if ! systemctl is-active --quiet ssh.socket &&
+        ! systemctl is-active --quiet sshd.socket; then
+      systemctl start "$ssh_service" || return 1
+    else
+      return 0
+    fi
+  fi
+  systemctl reload-or-restart "$ssh_service"
+}
 
-  # Reload or restart ssh service
-  systemctl reload-or-restart "$ssh_service" || exit $ERR_CONFIG_PUBKEY_AUTH_SSH
+disableSSHPubkeyAuth() {
+  local policy_dir="${SSH_POLICY_DIR:-/etc/aks/ssh-policy}"
+  local original
+  mkdir -p "$policy_dir" || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+  chmod 0700 "$policy_dir" || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+  if [ ! -f "$policy_dir/pubkey-auth" ]; then
+    install -d -m 0755 -o root -g root /run/sshd || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+    original=$(sshd -T -f "${SSHD_CONFIG_FILE:-/etc/ssh/sshd_config}") || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+    original=$(awk 'tolower($1) == "pubkeyauthentication" { print $2 }' <<< "$original")
+    case "$original" in
+      yes|no) ;;
+      *) echo "Cannot determine the original SSH public-key policy" >&2; exit "$ERR_CONFIG_PUBKEY_AUTH_SSH" ;;
+    esac
+    printf '%s\n' "$original" > "$policy_dir/pubkey-auth" || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+  fi
+  configureSSHPubkeyAuth no || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+  restoreSSH || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
+  reloadSSHAuthentication || exit "$ERR_CONFIG_PUBKEY_AUTH_SSH"
 }
 
 get_compute_sku() {
diff --git a/parts/linux/cloud-init/artifacts/cse_main.sh b/parts/linux/cloud-init/artifacts/cse_main.sh
index 0b6bc293576..d26d949d143 100755
--- a/parts/linux/cloud-init/artifacts/cse_main.sh
+++ b/parts/linux/cloud-init/artifacts/cse_main.sh
@@ -434,12 +434,7 @@ EOF
 # IMPORTANT: This stage should only run when actually joining a node to the cluster. This step should not be run when creating a VHD image
 function nodePrep {
     # Apply node access policy here so PIS images retain the distribution's SSH defaults.
-    if [ "${DISABLE_SSH}" = "true" ]; then
-        disableSSH || exit "$ERR_DISABLE_SSH"
-    elif [ "${DISABLE_PUBKEY_AUTH}" = "true" ]; then
-        logs_to_events "AKS.CSE.disableSSHPubkeyAuth" disableSSHPubkeyAuth
-    fi
-
+    logs_to_events "AKS.CSE.applySSHPolicy" applySSHPolicy || exit "$?"
     logs_to_events "AKS.CSE.configureAzureJson" configureAzureJson
     logs_to_events "AKS.CSE.ensureKubeCACert" ensureKubeCACert
 
diff --git a/pkg/agent/baker.go b/pkg/agent/baker.go
index 8f10356fa86..ff4188131df 100644
--- a/pkg/agent/baker.go
+++ b/pkg/agent/baker.go
@@ -1663,6 +1663,9 @@ func getContainerServiceFuncMap(config *datamodel.NodeBootstrappingConfiguration
 		"ShouldTurnOffPubkeyAuthSSH": func() bool {
 			return config.SSHStatus == datamodel.EntraIDSSH
 		},
+		"ShouldRestoreSSH": func() bool {
+			return config.SSHStatus == datamodel.SSHOn || config.SSHStatus == datamodel.EntraIDSSH
+		},
 		"GetSysctlContent": func() (string, error) {
 			templateFuncMap := make(template.FuncMap)
 			templateFuncMap["getPortRangeEndValue"] = getPortRangeEndValue
diff --git a/pkg/agent/baker_test.go b/pkg/agent/baker_test.go
index d61c448e979..d7662c75961 100644
--- a/pkg/agent/baker_test.go
+++ b/pkg/agent/baker_test.go
@@ -1092,6 +1092,20 @@ var _ = Describe("getLinuxNodeCSECommand", func() {
 		vars := decodeCSEVars(cseCmd)
 		Expect(vars).To(HaveKeyWithValue("DISABLE_SSH", "true"))
 		Expect(vars).To(HaveKeyWithValue("DISABLE_PUBKEY_AUTH", "false"))
+		Expect(vars).To(HaveKeyWithValue("SHOULD_RESTORE_SSH", "false"))
+	})
+
+	It("should restore SSH only for an explicit enabled policy", func() {
+		for _, status := range []datamodel.SSHStatus{datamodel.SSHUnspecified, datamodel.SSHOn, datamodel.EntraIDSSH} {
+			baseConfig.SSHStatus = status
+			vars := decodeCSEVars(templateGenerator.getLinuxNodeCSECommand(baseConfig))
+			expected := "true"
+			if status == datamodel.SSHUnspecified {
+				expected = "false"
+			}
+			Expect(vars).To(HaveKeyWithValue("SHOULD_RESTORE_SSH", expected))
+			Expect(vars).To(HaveKeyWithValue("DISABLE_SSH", "false"))
+		}
 	})
 
 	It("should set CUSTOM_DATA_ONLY_PROVISIONING", func() {
diff --git a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh
index bbf7cb3011b..1179148b00a 100755
--- a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh
+++ b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh
@@ -380,8 +380,6 @@ $cpu_pressure_commands"
             TEST_EVENTS_DIR="$TEST_EVENTS_ROOT/events"
             mkdir "$TEST_EVENTS_DIR"
             EVENTS_LOGGING_DIR=$TEST_EVENTS_DIR
-            SSH_POLICY_DIR="$TEST_EVENTS_ROOT/policy"
-            SSH_SYSTEMD_UNIT_DIR="$TEST_EVENTS_ROOT/systemd"
         }
 
         cleanup_ssh_telemetry() {
@@ -395,9 +393,21 @@ $cpu_pressure_commands"
         BeforeEach setup_ssh_telemetry
         AfterEach cleanup_ssh_telemetry
 
+        installSSHDisablePolicy() { :; }
+
         systemctl() {
             case "$1" in
-                daemon-reload) [ "$*" != "${FAIL_COMMAND:-}" ] ;;
+                daemon-reload) return 0 ;;
+                show)
+                    if [ "$3" = --property=ActiveState ]; then
+                        echo inactive
+                    else
+                        case " ${MISSING_UNITS:-} " in
+                            *" $2 "*) echo not-found ;;
+                            *) echo loaded ;;
+                        esac
+                    fi
+                    ;;
                 cat)
                     # MISSING_UNITS simulates a distro where the unit isn't installed.
                     case " ${MISSING_UNITS:-} " in
@@ -422,12 +432,16 @@ $cpu_pressure_commands"
             :
         }
 
-        It 'stops every SSH unit without changing its boot activation'
+        It 'stops and disables every SSH unit that is present'
             When call disableSSH
             The output should equal "stop ssh.socket
+disable ssh.socket
 stop sshd.socket
+disable sshd.socket
 stop ssh
-stop sshd"
+disable ssh
+stop sshd
+disable sshd"
             The status should be success
             SSH_EVENT_FILTER='
                 length == 10 and
@@ -439,7 +453,7 @@ stop sshd"
                 ([.[] | select(.TaskName == "AKS.CSE.disableSSH.unit") | .Message | fromjson |
                   select(.Phase == "Completed")] |
                     length == 4 and all(.[];
-                        .CatExitCode == 0 and .StopExitCode == 0 and .DisableExitCode == null and
+                        .CatExitCode == 0 and .StopExitCode == 0 and .DisableExitCode == 0 and
                         .Skipped == false and .SchemaVersion == 1)) and
                     all(.[]; (.Timestamp | length) > 0 and (.OperationId | length) > 0) and
                 ([.[].Message | fromjson | select(.Phase == "Started")] |
@@ -449,34 +463,38 @@ stop sshd"
         End
 
         It 'skips units that are not installed'
-            MISSING_UNITS="ssh.socket ssh sshd"
+            MISSING_UNITS="ssh sshd"
             When call disableSSH
-            The output should equal "stop sshd.socket"
+            The output should equal "stop ssh.socket
+disable ssh.socket
+stop sshd.socket
+disable sshd.socket"
             The status should be success
             SSH_EVENT_FILTER='
                 [.[].Message | fromjson | select(.Phase == "Completed" and .Skipped == true)] |
-                length == 3 and all(.[]; .CatExitCode == 1 and .StopExitCode == null and
+                length == 2 and all(.[]; .CatExitCode == 1 and .StopExitCode == null and
                     .DisableExitCode == null and .ExitCode == 0)
             '
             The result of function ssh_events_match should equal true
         End
 
         It 'skips sshd.socket on distros without socket activation'
-            MISSING_UNITS="ssh.socket ssh sshd.socket"
+            MISSING_UNITS="ssh ssh.socket sshd.socket"
             When call disableSSH
-            The output should equal "stop sshd"
+            The output should equal "stop sshd
+disable sshd"
             The status should be success
         End
 
-        It 'stops listeners and fails provisioning when the persistent policy cannot be loaded'
-            FAIL_COMMAND="daemon-reload"
+        It 'exits with the CSE SSH error code when a unit cannot be disabled'
+            FAIL_COMMAND="disable sshd.socket"
             When run disableSSH
-            The stderr should include "Failed to install persistent SSH disable policy"
-            The output should include "stop sshd.socket"
+            The output should include "sshd.socket could not be disabled"
             The status should equal 172
             SSH_EVENT_FILTER='
                 [.[] | select(.EventLevel == "Error") | .Message | fromjson] |
-                length == 1 and
+                length == 2 and
+                any(.[]; .Unit == "sshd.socket" and .StopExitCode == 0 and .DisableExitCode == 1) and
                 any(.[]; .Unit == "" and .ExitCode == 172 and .Phase == "Completed")
             '
             The result of function ssh_events_match should equal true
@@ -486,14 +504,14 @@ stop sshd"
             FAIL_COMMAND="stop sshd"
             When run disableSSH
             The output should include "sshd could not be stopped"
-            The output should include "stop sshd"
-            The output should include "stop sshd.socket"
+            The output should include "disable sshd"
+            The output should include "disable sshd.socket"
             The status should equal 172
             SSH_EVENT_FILTER='
                 [.[].Message | fromjson] |
-                any(.[]; .Unit == "sshd" and .StopExitCode == 1 and .DisableExitCode == null) and
+                any(.[]; .Unit == "sshd" and .StopExitCode == 1 and .DisableExitCode == 0) and
                 any(.[]; .Unit == "" and .ExitCode == 172 and .Phase == "Completed") and
-                any(.[]; .Unit == "sshd.socket" and .StopExitCode == 0 and .DisableExitCode == null)
+                any(.[]; .Unit == "sshd.socket" and .StopExitCode == 0 and .DisableExitCode == 0)
             '
             The result of function ssh_events_match should equal true
         End
@@ -501,7 +519,7 @@ stop sshd"
         It 'emits only the compact outcome schema'
             When call disableSSH
             The status should be success
-            The output should include "stop sshd.socket"
+            The output should include "disable sshd.socket"
             SSH_EVENT_FILTER='
                 [.[].Message | fromjson | select(.Phase == "Completed" and .Unit != "") | keys] |
                 length == 4 and all(.[]; . ==
@@ -518,7 +536,7 @@ stop sshd"
             }
             When call disable_ssh_twice
             The status should be success
-            The output should include "stop sshd.socket"
+            The output should include "disable sshd.socket"
             SSH_EVENT_FILTER='
                 length == 20 and ([.[].Message | fromjson | .AttemptId] | unique | length == 2)
             '
@@ -540,7 +558,7 @@ stop sshd"
 
             When call disableSSH
             The status should be success
-            The output should include "stop sshd.socket"
+            The output should include "disable sshd.socket"
             The stderr should equal ""
             The file "$TEST_EVENTS_ROOT/cleaned" should be exist
             SSH_EVENT_FILTER='length == 10 and all(.[]; (.Message | fromjson | .SchemaVersion) == 1)'
@@ -565,7 +583,7 @@ stop sshd"
             EVENTS_LOGGING_DIR="$TEST_EVENTS_DIR/not-a-directory"
             When call disableSSH
             The status should be success
-            The output should include "stop sshd.socket"
+            The output should include "disable sshd.socket"
             The stderr should include "WARNING: could not emit SSH disable"
         End
 
@@ -573,7 +591,7 @@ stop sshd"
             jq() { return 1; }
             When call disableSSH
             The status should be success
-            The output should include "stop sshd.socket"
+            The output should include "disable sshd.socket"
             The stderr should include "WARNING: could not emit SSH disable"
             The stderr should include "WARNING: could not serialize SSH disable result"
         End
@@ -584,7 +602,7 @@ stop sshd"
             When run disableSSH
             The status should equal 172
             The output should include "sshd could not be stopped"
-            The output should include "stop sshd"
+            The output should include "disable sshd"
             The stderr should include "WARNING: could not emit SSH disable"
             The stderr should include "WARNING: could not serialize SSH disable result"
         End
@@ -592,15 +610,22 @@ stop sshd"
     Describe 'disableSSHPubkeyAuth'
         setup() {
             SSHD_CONFIG_FILE="$(mktemp)"
+            SSH_POLICY_DIR="$(mktemp -d)"
             SSH_SERVICE_ACTIVE="true"
             SSH_SERVICE_EXISTS="true"
+            SSHD_OUTPUT_KEY=pubkeyauthentication
+            DISABLE_SSH=false
+            DISABLE_PUBKEY_AUTH=false
+            SHOULD_RESTORE_SSH=false
         }
 
         cleanup() {
             rm -f "${SSHD_CONFIG_FILE}"
+            rm -rf "${SSH_POLICY_DIR}"
             unset SSHD_CONFIG_FILE
             unset SSH_SERVICE_ACTIVE
             unset SSH_SERVICE_EXISTS
+            unset SSHD_OUTPUT_KEY
         }
 
         systemctl() {
@@ -608,18 +633,26 @@ stop sshd"
             if [ "$1" = "cat" ] && [ "${SSH_SERVICE_EXISTS}" != "true" ]; then
                 return 1
             fi
-            if [ "$1" = "is-active" ] && [ "${SSH_SERVICE_ACTIVE}" != "true" ]; then
-                return 3
+            if [ "$1" = "is-active" ]; then
+                case "$3" in
+                    *.socket) return 3 ;;
+                esac
+                [ "${SSH_SERVICE_ACTIVE}" = "true" ] || return 3
             fi
             return 0
         }
 
         sshd() {
             echo "sshd $*" >&2
+            if [ "$1" = -T ]; then
+                awk -v key="$SSHD_OUTPUT_KEY" 'BEGIN { value="yes" } /^Match/ { exit } /^PubkeyAuthentication / { value=$2 }
+                    END { print key " " value }' "$3"
+            fi
             return 0
         }
 
         install() {
+            [ "$1" != -d ] || return 0
             while [ "$#" -gt 2 ]; do
                 shift
             done
@@ -631,6 +664,8 @@ stop sshd"
             cat "${SSHD_CONFIG_FILE}"
         }
 
+        restoreSSH() { return 0; }
+
         BeforeEach 'setup'
         AfterEach 'cleanup'
 
@@ -709,6 +744,599 @@ PubkeyAuthentication no"
 PubkeyAuthentication no"
             The stderr should include "systemctl start sshd.service"
         End
+
+        It 'does not start SSH when candidate validation fails'
+            echo "PubkeyAuthentication yes" > "${SSHD_CONFIG_FILE}"
+            SSH_SERVICE_ACTIVE=false
+            sshd() {
+                if [ "$1" = -T ]; then
+                    echo "pubkeyauthentication yes"
+                    return 0
+                fi
+                return 1
+            }
+            When run disableSSHPubkeyAuth
+            The status should equal 178
+            The contents of file "${SSHD_CONFIG_FILE}" should equal "PubkeyAuthentication yes"
+            The stderr should not include "systemctl start"
+        End
+
+        It 'installs the restrictive config before starting SSH'
+            echo "PubkeyAuthentication yes" > "${SSHD_CONFIG_FILE}"
+            SSH_SERVICE_ACTIVE=false
+            systemctl() {
+                case "$1" in
+                    cat) return 0 ;;
+                    is-active) return 3 ;;
+                    start) grep -qx 'PubkeyAuthentication no' "${SSHD_CONFIG_FILE}" ;;
+                    reload-or-restart) return 0 ;;
+                    *) return 1 ;;
+                esac
+            }
+            When call disableSSHPubkeyAuth
+            The status should be success
+            The stderr should include "sshd -t -f"
+        End
+
+        Describe 'authentication restoration across sshd output formats'
+            Parameters
+                pubkeyauthentication
+                PubkeyAuthentication
+                PUBKEYAUTHENTICATION
+            End
+
+            It "restores only the owned global setting with $1 in sshd output"
+                SSHD_OUTPUT_KEY="$1"
+                cat > "${SSHD_CONFIG_FILE}" <<'EOF'
+PubkeyAuthentication yes
+Match User entra
+  PubkeyAuthentication yes
+EOF
+                disable_then_restore_auth() {
+                    disableSSHPubkeyAuth
+                    SHOULD_RESTORE_SSH=true
+                    applySSHPolicy
+                    cat "${SSHD_CONFIG_FILE}"
+                }
+                When call disable_then_restore_auth
+                The status should be success
+                The output should equal "PubkeyAuthentication yes
+Match User entra
+  PubkeyAuthentication yes"
+                The stderr should include "systemctl reload-or-restart ssh.service"
+                The file "$SSH_POLICY_DIR/pubkey-auth" should not be exist
+            End
+        End
+
+        It 'does not activate SSH when installing the authentication config fails'
+            echo "PubkeyAuthentication yes" > "${SSHD_CONFIG_FILE}"
+            SSH_SERVICE_ACTIVE=false
+            install() { [ "$1" = -d ]; }
+            When run disableSSHPubkeyAuth
+            The status should equal 178
+            The contents of file "${SSHD_CONFIG_FILE}" should equal "PubkeyAuthentication yes"
+            The stderr should include "sshd -t -f"
+            The stderr should not include "systemctl start"
+        End
+
+        It 'does not reopen SSH when an Include overrides the requested authentication policy'
+            echo "Include /etc/ssh/sshd_config.d/*.conf" > "$SSHD_CONFIG_FILE"
+            SSH_SERVICE_ACTIVE=false
+            sshd() {
+                [ "$1" != -T ] || echo 'pubkeyauthentication yes'
+                return 0
+            }
+            When run disableSSHPubkeyAuth
+            The status should equal 178
+            The stderr should include "does not enforce the requested public-key policy"
+            The stderr should not include "systemctl start"
+            The contents of file "$SSHD_CONFIG_FILE" should equal "Include /etc/ssh/sshd_config.d/*.conf"
+        End
+    End
+    Describe 'persistent SSH access policy'
+        Include "./parts/linux/cloud-init/artifacts/cse_config.sh"
+        Include "./spec/support/ssh_policy.sh"
+        BeforeEach setup_ssh_policy
+        AfterEach cleanup_ssh_policy
+
+        run_node_prep_ssh_reconciliation() {
+            sed -n '/^function nodePrep {/,/^}/p' \
+                parts/linux/cloud-init/artifacts/cse_main.sh |
+                sed "s|/opt/azure/containers/base_prep.complete|$SSH_TEST_ROOT/base_prep.complete|g" \
+                    > "$SSH_TEST_ROOT/node-prep.sh"
+            # Load the actual entry point without executing CSE's top-level setup.
+            . "$SSH_TEST_ROOT/node-prep.sh"
+            nodePrep
+        }
+
+        It 'blocks a dependency start even after a package unmask'
+            disable_and_unmask() {
+                disableSSH
+                systemctl unmask sshd.service
+                systemctl start sshd.service
+                cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+            }
+            When call disable_and_unmask
+            The status should be success
+            The output should equal inactive
+        End
+
+        It 'does not lose the original activation mode on repeated disable'
+            disable_twice_and_restore() {
+                disableSSH
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+                cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                cat "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            }
+            When call disable_twice_and_restore
+            The status should be success
+            The output should equal activeenabled
+            The file "$SSH_POLICY_DIR/disabled" should not be exist
+        End
+
+        It 'can restore a unique endpoint disabled by legacy CSE without inventing boot links'
+            restore_legacy_disabled() {
+                printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+                cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                cat "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            }
+            When call restore_legacy_disabled
+            The status should be success
+            The output should equal activedisabled
+        End
+
+        It 'explicitly enables the unique legacy service when both endpoints are disabled'
+            restore_ambiguous_legacy() {
+                printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                printf loaded > "$SSH_MODEL_DIR/sshd.socket.LoadState"
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+                cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                cat "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                cat "$SSH_MODEL_DIR/sshd.socket.UnitFileState"
+            }
+            When call restore_ambiguous_legacy
+            The status should be success
+            The output should equal activeenableddisabled
+            The file "$SSH_POLICY_DIR/disabled" should not be exist
+        End
+
+        It 'leaves a stopped service unchanged for LocalUser without an owned policy'
+            DISABLE_SSH=false
+            SHOULD_RESTORE_SSH=true
+            printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+            When call applySSHPolicy
+            The status should be success
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+            The contents of file "$SSH_MODEL_DIR/sshd.service.UnitFileState" should equal disabled
+            The file "$SSH_TEST_ROOT/activation-commands" should not be exist
+        End
+
+        Describe 'legacy Entra authentication reconciliation without a disable marker'
+            prepare_legacy_entra_auth() {
+                DISABLE_SSH=false
+                DISABLE_PUBKEY_AUTH=true
+                SHOULD_RESTORE_SSH=true
+                EXPECTED_PUBKEY=no
+                printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+            }
+
+            It 'enables boot activation after validating PubkeyAuthentication=no'
+                prepare_legacy_entra_auth
+                When call applySSHPolicy
+                The status should be success
+                The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal active
+                The contents of file "$SSH_MODEL_DIR/sshd.service.UnitFileState" should equal enabled
+                The contents of file "$SSHD_CONFIG_FILE" should include "PubkeyAuthentication no"
+                The file "$SSH_POLICY_DIR/disabled" should not be exist
+            End
+
+            It 'preserves an external service mask with PubkeyAuthentication=no'
+                prepare_legacy_entra_auth
+                printf masked > "$SSH_MODEL_DIR/sshd.service.LoadState"
+                When run applySSHPolicy
+                The status should equal 178
+                The stderr should include "external masks are not removed"
+                The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+                The contents of file "$SSH_MODEL_DIR/sshd.service.LoadState" should equal masked
+                The contents of file "$SSH_MODEL_DIR/sshd.service.UnitFileState" should equal disabled
+            End
+        End
+
+        Describe 'LocalUser authentication restoration without a disable marker'
+            Parameters
+                loaded active
+                loaded inactive
+                masked inactive
+            End
+
+            It "restores authentication without changing a $1/$2 service"
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                mkdir -p "$SSH_POLICY_DIR"
+                printf yes > "$SSH_POLICY_DIR/pubkey-auth"
+                printf 'PubkeyAuthentication no\nMatch User entra\n  PubkeyAuthentication yes\n' > "$SSHD_CONFIG_FILE"
+                printf '%s' "$1" > "$SSH_MODEL_DIR/sshd.service.LoadState"
+                printf '%s' "$2" > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                When call applySSHPolicy
+                The status should be success
+                The contents of file "$SSHD_CONFIG_FILE" should equal "PubkeyAuthentication yes
+Match User entra
+  PubkeyAuthentication yes"
+                The contents of file "$SSH_MODEL_DIR/sshd.service.LoadState" should equal "$1"
+                The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal "$2"
+                The contents of file "$SSH_MODEL_DIR/sshd.service.UnitFileState" should equal disabled
+                The file "$SSH_TEST_ROOT/activation-commands" should not be exist
+                The file "$SSH_POLICY_DIR/pubkey-auth" should not be exist
+            End
+        End
+
+        It 'keeps the authentication backup when reloading LocalUser authentication fails'
+            DISABLE_SSH=false
+            SHOULD_RESTORE_SSH=true
+            mkdir -p "$SSH_POLICY_DIR"
+            printf yes > "$SSH_POLICY_DIR/pubkey-auth"
+            printf 'PubkeyAuthentication no\n' > "$SSHD_CONFIG_FILE"
+            FAIL_COMMAND="reload-or-restart sshd.service"
+            When call applySSHPolicy
+            The status should equal 178
+            The contents of file "$SSH_POLICY_DIR/pubkey-auth" should equal yes
+            The file "$SSH_TEST_ROOT/activation-commands" should not be exist
+        End
+
+        It 'restores authentication before activating an owned disabled service for LocalUser'
+            restore_owned_localuser_auth() {
+                disableSSHPubkeyAuth
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                EXPECTED_PUBKEY=yes
+                applySSHPolicy
+            }
+            When call restore_owned_localuser_auth
+            The status should be success
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal active
+            The contents of file "$SSH_MODEL_DIR/sshd.service.UnitFileState" should equal enabled
+            The contents of file "$SSHD_CONFIG_FILE" should include "PubkeyAuthentication yes"
+            The file "$SSH_POLICY_DIR/disabled" should not be exist
+            The file "$SSH_POLICY_DIR/pubkey-auth" should not be exist
+        End
+
+        It 'deduplicates Ubuntu aliases when selecting the legacy service'
+            restore_legacy_ubuntu() {
+                SSH_MODEL_OS=ubuntu
+                printf loaded > "$SSH_MODEL_DIR/ssh.service.LoadState"
+                printf loaded > "$SSH_MODEL_DIR/ssh.socket.LoadState"
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+                cat "$SSH_MODEL_DIR/ssh.service.ActiveState"
+                cat "$SSH_MODEL_DIR/ssh.service.UnitFileState"
+                cat "$SSH_MODEL_DIR/ssh.socket.UnitFileState"
+            }
+            When call restore_legacy_ubuntu
+            The status should be success
+            The output should equal activeenableddisabled
+        End
+
+        It 'refuses a legacy restore with multiple distinct services'
+            restore_distinct_legacy_services() {
+                printf loaded > "$SSH_MODEL_DIR/ssh.service.LoadState"
+                printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+            }
+            When call restore_distinct_legacy_services
+            The status should equal 172
+            The stderr should include "multiple distinct services"
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+        End
+
+        It 'does not bypass an external service mask in the legacy fallback'
+            restore_masked_legacy_service() {
+                printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                printf loaded > "$SSH_MODEL_DIR/sshd.socket.LoadState"
+                disableSSH
+                printf masked > "$SSH_MODEL_DIR/sshd.service.LoadState"
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+            }
+            When call restore_masked_legacy_service
+            The status should equal 172
+            The stderr should include "external masks are not removed"
+            The contents of file "$SSH_MODEL_DIR/sshd.service.LoadState" should equal masked
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+        End
+
+        It 'keeps SSH blocked when the enabled policy is unspecified'
+            disable_then_unspecified() {
+                disableSSH
+                DISABLE_SSH=false
+                applySSHPolicy
+                systemctl start sshd.service
+                cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+            }
+            When call disable_then_unspecified
+            The status should be success
+            The output should equal inactive
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+        End
+
+        It 'stops an already running masked service without removing its mask'
+            printf masked > "$SSH_MODEL_DIR/sshd.service.LoadState"
+            When call disableSSH
+            The status should be success
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+            The contents of file "$SSH_MODEL_DIR/sshd.service.LoadState" should equal masked
+        End
+
+        It 'reports an external mask conflict without reopening SSH'
+            restore_externally_masked() {
+                disableSSH
+                printf masked > "$SSH_MODEL_DIR/sshd.service.LoadState"
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+            }
+            When call restore_externally_masked
+            The status should equal 172
+            The stderr should include "external masks are not removed"
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+            The contents of file "$SSH_MODEL_DIR/sshd.service.LoadState" should equal masked
+        End
+
+        It 'preserves an administrator drop-in and fails provisioning'
+            mkdir -p "$SSH_SYSTEMD_UNIT_DIR/sshd.service.d"
+            echo '[Unit]' > "$SSH_SYSTEMD_UNIT_DIR/sshd.service.d/90-aks-ssh-access.conf"
+            When run disableSSH
+            The status should equal 172
+            The stderr should include "Refusing to overwrite"
+            The contents of file "$SSH_SYSTEMD_UNIT_DIR/sshd.service.d/90-aks-ssh-access.conf" should equal '[Unit]'
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+        End
+
+        It 'leaves an external mask unchanged for LocalUser without an owned policy'
+            DISABLE_SSH=false
+            SHOULD_RESTORE_SSH=true
+            printf masked > "$SSH_MODEL_DIR/sshd.service.LoadState"
+            When call applySSHPolicy
+            The status should be success
+            The contents of file "$SSH_MODEL_DIR/sshd.service.LoadState" should equal masked
+            The file "$SSH_TEST_ROOT/activation-commands" should not be exist
+            The file "$SSH_POLICY_DIR/disabled" should not be exist
+        End
+
+        It 'does not release a disable marker when its ownership record is missing'
+            restore_with_missing_record() {
+                disableSSH
+                rm "$SSH_POLICY_DIR/activation"
+                printf enabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+            }
+            When call restore_with_missing_record
+            The status should equal 172
+            The stderr should include "owned activation record is missing"
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+        End
+
+        It 'fails provisioning when the persistent policy cannot be loaded'
+            FAIL_COMMAND=daemon-reload
+            When run disableSSH
+            The status should equal 172
+            The stderr should include "Failed to install persistent SSH disable policy"
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+        End
+
+        It 'restores socket activation from a disabled PIS cache without enabling its service'
+            restore_ubuntu_socket() {
+                SSH_MODEL_OS=ubuntu
+                printf loaded > "$SSH_MODEL_DIR/ssh.service.LoadState"
+                printf loaded > "$SSH_MODEL_DIR/ssh.socket.LoadState"
+                printf enabled > "$SSH_MODEL_DIR/ssh.socket.UnitFileState"
+                printf active > "$SSH_MODEL_DIR/ssh.socket.ActiveState"
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+                cat "$SSH_MODEL_DIR/ssh.socket.ActiveState"
+                cat "$SSH_MODEL_DIR/ssh.service.UnitFileState"
+            }
+            When call restore_ubuntu_socket
+            The status should be success
+            The output should equal activedisabled
+        End
+
+        It 'respects a package migration to a newly installed socket'
+            restore_migrated_socket() {
+                disableSSH
+                printf loaded > "$SSH_MODEL_DIR/sshd.socket.LoadState"
+                printf enabled > "$SSH_MODEL_DIR/sshd.socket.UnitFileState"
+                systemctl start sshd.socket
+                [ "$(cat "$SSH_MODEL_DIR/sshd.socket.ActiveState")" = inactive ] || return 1
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+                cat "$SSH_MODEL_DIR/sshd.socket.ActiveState"
+                cat "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            }
+            When call restore_migrated_socket
+            The status should be success
+            The output should equal activedisabled
+        End
+
+        It 'reinstates the persistent block if activation fails'
+            fail_restore() {
+                disableSSH
+                FAIL_COMMAND="start sshd.service"
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                applySSHPolicy
+            }
+            When call fail_restore
+            The status should equal 172
+            The stderr should include "Failed to restore SSH activation"
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+        End
+
+        It 'configures Entra authentication before restoring a disabled service'
+            restore_entra_auth() {
+                disableSSH
+                DISABLE_SSH=false
+                DISABLE_PUBKEY_AUTH=true
+                EXPECTED_PUBKEY=no
+                applySSHPolicy
+                cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+            }
+            When call restore_entra_auth
+            The status should be success
+            The output should equal active
+            The contents of file "$SSHD_CONFIG_FILE" should include "PubkeyAuthentication no"
+            The contents of file "$SSHD_CONFIG_FILE" should include "  PubkeyAuthentication yes"
+        End
+
+        Describe 'owned Entra activation failures'
+            Parameters
+                syntax
+                install
+            End
+
+            It "keeps the owned block when authentication $1 fails"
+                SSH_AUTH_FAILURE="$1"
+                sshd() {
+                    if [ "$1" = -T ]; then
+                        awk 'BEGIN { value="yes" } /^Match/ { exit } /^PubkeyAuthentication / { value=$2 }
+                            END { print "pubkeyauthentication " value }' "$3"
+                    else
+                        [ "$SSH_AUTH_FAILURE" != syntax ]
+                    fi
+                }
+                install() {
+                    [ "$1" != -d ] || return 0
+                    [ "$SSH_AUTH_FAILURE" != install ] || return 1
+                    while [ "$#" -gt 2 ]; do shift; done
+                    command cp "$1" "$2"
+                }
+                fail_owned_entra() {
+                    disableSSH
+                    DISABLE_SSH=false
+                    DISABLE_PUBKEY_AUTH=true
+                    applySSHPolicy
+                }
+                When run fail_owned_entra
+                The status should equal 178
+                The file "$SSH_POLICY_DIR/disabled" should be exist
+                The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+                The contents of file "$SSHD_CONFIG_FILE" should include "PubkeyAuthentication yes"
+            End
+        End
+
+        It 'gives SSH Disabled precedence over the authentication request'
+            DISABLE_PUBKEY_AUTH=true
+            When call applySSHPolicy
+            The status should be success
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+            The file "$SSH_POLICY_DIR/pubkey-auth" should not be exist
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+        End
+
+        It 'enforces Disabled at nodePrep even when basePrep was cached as enabled'
+            touch "$SSH_TEST_ROOT/base_prep.complete"
+            logs_to_events() {
+                [ "$2" = applySSHPolicy ] || exit 1
+                "$2"
+                exit "$?"
+            }
+            When run run_node_prep_ssh_reconciliation
+            The status should be success
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
+        End
+
+        It 'restores explicitly enabled SSH at nodePrep from a disabled cache'
+            prepare_disabled_cache() {
+                touch "$SSH_TEST_ROOT/base_prep.complete"
+                disableSSH
+                DISABLE_SSH=false
+                SHOULD_RESTORE_SSH=true
+                run_node_prep_ssh_reconciliation
+            }
+            logs_to_events() {
+                [ "$2" = applySSHPolicy ] || exit 1
+                "$2"
+                exit "$?"
+            }
+            When run prepare_disabled_cache
+            The status should be success
+            The file "$SSH_POLICY_DIR/disabled" should not be exist
+            The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal active
+        End
+
+        Describe 'uncached Entra node preparation'
+            Parameters
+                legacy active
+                socket inactive
+            End
+
+            It "does not activate a stopped service twice after restoring a $1 snapshot"
+                prepare_uncached_entra() {
+                    printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                    printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+                    if [ "$1" = socket ]; then
+                        printf loaded > "$SSH_MODEL_DIR/sshd.socket.LoadState"
+                        printf enabled > "$SSH_MODEL_DIR/sshd.socket.UnitFileState"
+                        printf active > "$SSH_MODEL_DIR/sshd.socket.ActiveState"
+                    fi
+                    disableSSH
+                    DISABLE_SSH=false
+                    DISABLE_PUBKEY_AUTH=true
+                    applySSHPolicy
+                    [ ! -f "$SSH_POLICY_DIR/disabled" ] || return 1
+                    [ "$(cat "$SSH_MODEL_DIR/sshd.service.ActiveState")" = "$2" ] || return 1
+                    run_node_prep_ssh_reconciliation
+                }
+                logs_to_events() {
+                    case "$2" in
+                        applySSHPolicy|disableSSHPubkeyAuth)
+                            printf '%s\n' "$2" >> "$SSH_TEST_ROOT/policy-events"
+                            "$2"
+                            ;;
+                        configureAzureJson) exit 0 ;;
+                        *) exit 1 ;;
+                    esac
+                }
+                When run prepare_uncached_entra "$1" "$2"
+                The status should be success
+                The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal "$2"
+                The contents of file "$SSH_TEST_ROOT/commands" should not include 'enable sshd.service'
+                The contents of file "$SSH_TEST_ROOT/policy-events" should equal "disableSSHPubkeyAuth
+applySSHPolicy
+disableSSHPubkeyAuth"
+            End
+        End
     End
     Describe 'configureAzureJson'
         AZURE_JSON_PATH="azure.json"
diff --git a/spec/parts/linux/cloud-init/artifacts/cse_main_pis_reconcile_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_main_pis_reconcile_spec.sh
index ce732a86c91..97a7589c1dc 100644
--- a/spec/parts/linux/cloud-init/artifacts/cse_main_pis_reconcile_spec.sh
+++ b/spec/parts/linux/cloud-init/artifacts/cse_main_pis_reconcile_spec.sh
@@ -82,7 +82,7 @@ Describe 'cse_main.sh PIS-safe configuration'
 
         It 'applies live SSH policy before writing azure.json and ca.crt'
             When call node_prep_event_calls
-            The line 1 of output should equal "disableSSHPubkeyAuth"
+            The line 1 of output should equal "applySSHPolicy"
             The line 2 of output should equal "configureAzureJson"
             The line 3 of output should equal "ensureKubeCACert"
         End
@@ -130,7 +130,7 @@ Describe 'cse_main.sh PIS-safe configuration'
     End
 
     It 'does not apply SSH access policy during image preparation'
-        When call phase_count basePrep disableSSH
+        When call phase_count basePrep applySSHPolicy
         The output should equal "0"
     End
 
@@ -138,6 +138,8 @@ Describe 'cse_main.sh PIS-safe configuration'
         prepare_node_prep() {
             NODE_PREP_SCRIPT=$(mktemp)
             phase_body nodePrep > "$NODE_PREP_SCRIPT"
+            sed -n '/^applySSHPolicy() {/,/^}/p' \
+                ./parts/linux/cloud-init/artifacts/cse_config.sh >> "$NODE_PREP_SCRIPT"
         }
         cleanup_node_prep() { rm -f "$NODE_PREP_SCRIPT"; }
         BeforeEach prepare_node_prep
diff --git a/spec/parts/linux/cloud-init/artifacts/cse_ssh_policy_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_ssh_policy_spec.sh
index cbf572fc39d..d6f80fa7366 100644
--- a/spec/parts/linux/cloud-init/artifacts/cse_ssh_policy_spec.sh
+++ b/spec/parts/linux/cloud-init/artifacts/cse_ssh_policy_spec.sh
@@ -1,155 +1,181 @@
 #!/bin/bash
 
-Describe 'persistent SSH disablement'
+Describe 'SSH activation policy helpers'
     CSE_CONFIG_GPU_FILEPATH="./parts/linux/cloud-init/artifacts/cse_config_gpu.sh"
     CSE_CONFIG_LOCALDNS_FILEPATH="./parts/linux/cloud-init/artifacts/cse_config_localdns.sh"
     CSE_CONFIG_KUBELET_FILEPATH="./parts/linux/cloud-init/artifacts/cse_config_kubelet.sh"
     CSE_CONFIG_NETWORK_FILEPATH="./parts/linux/cloud-init/artifacts/cse_config_network.sh"
     Include "./parts/linux/cloud-init/artifacts/cse_config.sh"
-    Include "./parts/linux/cloud-init/artifacts/cse_helpers.sh"
+    Include "./spec/support/ssh_policy.sh"
+    BeforeEach setup_ssh_policy
+    AfterEach cleanup_ssh_policy
 
-    setup() {
-        TEST_ROOT=$(mktemp -d)
-        SSH_POLICY_DIR="$TEST_ROOT/policy"
-        SSH_SYSTEMD_UNIT_DIR="$TEST_ROOT/systemd"
-        EVENTS_LOGGING_DIR="$TEST_ROOT/events"
-        MODEL="$TEST_ROOT/model"
-        mkdir -p "$MODEL"
+    save_stopped_policy() {
+        installSSHDisablePolicy || return 1
         local unit
-        for unit in ssh.service sshd.service ssh.socket sshd.socket; do
-            printf not-found > "$MODEL/$unit.LoadState"
-            printf disabled > "$MODEL/$unit.UnitFileState"
-            printf inactive > "$MODEL/$unit.ActiveState"
+        for unit in ssh.socket sshd.socket ssh.service sshd.service; do
+            systemctl stop "$unit" || return 1
+            systemctl disable "$unit" || return 1
         done
     }
 
-    cleanup() {
-        find "$TEST_ROOT" -type f -delete
-        find "$TEST_ROOT" -depth -type d -empty -delete
-    }
-
-    BeforeEach setup
-    AfterEach cleanup
+    It 'preserves the first activation snapshot and blocks dependency starts'
+        repeat_policy() {
+            save_stopped_policy
+            installSSHDisablePolicy
+            systemctl start sshd.service
+            cat "$SSH_POLICY_DIR/activation"
+            cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+        }
+        When call repeat_policy
+        The status should be success
+        The output should equal "sshd.service enabled
+inactive"
+    End
 
-    systemctl() {
-        local action="$1" unit="$2"
-        printf '%s\n' "$*" >> "$TEST_ROOT/commands"
-        [ "$action" != daemon-reload ] || return 0
-        case "$unit" in ssh|sshd) unit="$unit.service" ;; esac
-        case "$action" in
-            cat)
-                case "$(cat "$MODEL/$unit.LoadState")" in
-                    loaded|masked) return 0 ;;
-                    *) return 1 ;;
-                esac
-                ;;
-            stop)
-                # systemctl stop succeeds without clearing an existing failed state.
-                [ "$(cat "$MODEL/$unit.ActiveState")" = failed ] ||
-                    printf inactive > "$MODEL/$unit.ActiveState"
-                ;;
-            start)
-                [ "$(cat "$MODEL/$unit.LoadState")" = loaded ] || return 1
-                if [ -f "$SSH_POLICY_DIR/disabled" ] &&
-                    grep -Fxq "ConditionPathExists=!$SSH_POLICY_DIR/disabled" \
-                        "$SSH_SYSTEMD_UNIT_DIR/$unit.d/90-aks-ssh-access.conf"; then
-                    return 0
-                fi
-                printf active > "$MODEL/$unit.ActiveState"
-                ;;
-            *) echo "Unexpected systemctl command: $*" >&2; return 1 ;;
-        esac
-    }
+    It 'restores an owned service and leaves later LocalUser reconciliation unchanged'
+        restore_twice() {
+            save_stopped_policy
+            restoreSSH
+            DISABLE_SSH=false
+            SHOULD_RESTORE_SSH=true
+            applySSHPolicy
+            cat "$SSH_MODEL_DIR/sshd.service.ActiveState"
+            cat "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            grep -c '^start ' "$SSH_TEST_ROOT/activation-commands"
+        }
+        When call restore_twice
+        The status should be success
+        The output should equal "activeenabled1"
+        The file "$SSH_POLICY_DIR/disabled" should not be exist
+        The file "$SSH_POLICY_DIR/activation" should not be exist
+    End
 
-    timeout() { shift; "$@"; }
-    sleep() { :; }
+    It 'captures and restores an unambiguous stopped legacy service without boot links'
+        restore_legacy_snapshot() {
+            printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+            printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            save_stopped_policy
+            cat "$SSH_POLICY_DIR/activation"
+            restoreSSH
+        }
+        When call restore_legacy_snapshot
+        The status should be success
+        The output should equal 'sshd.service disabled'
+        The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal active
+        The contents of file "$SSH_MODEL_DIR/sshd.service.UnitFileState" should equal disabled
+        The contents of file "$SSH_TEST_ROOT/activation-commands" should equal 'start sshd.service'
+        The file "$SSH_POLICY_DIR/disabled" should not be exist
+        The file "$SSH_POLICY_DIR/activation" should not be exist
+    End
 
-    Describe 'service and socket starts'
+    Describe 'invalid ownership records'
         Parameters
-            ssh.service
-            sshd.service
-            ssh.socket
-            sshd.socket
+            missing 'owned activation record is missing'
+            invalid 'no valid saved activation mode'
         End
-        It "keeps $1 blocked after repeated disable without changing enablement"
-            block_start() {
-                printf loaded > "$MODEL/$1.LoadState"
-                printf enabled > "$MODEL/$1.UnitFileState"
-                printf active > "$MODEL/$1.ActiveState"
-                disableSSH
-                disableSSH
-                systemctl start "$1"
+
+        It "keeps its block when the saved record is $1"
+            break_record() {
+                save_stopped_policy
+                if [ "$1" = missing ]; then
+                    rm "$SSH_POLICY_DIR/activation"
+                else
+                    printf 'unknown enabled\n' > "$SSH_POLICY_DIR/activation"
+                fi
+                restoreSSH
             }
-            When call block_start "$1"
-            The status should be success
-            The contents of file "$MODEL/$1.ActiveState" should equal inactive
-            The contents of file "$MODEL/$1.UnitFileState" should equal enabled
-            The contents of file "$TEST_ROOT/commands" should not include 'disable '
+            When call break_record "$1"
+            The status should be failure
+            The stderr should include "$2"
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+            The file "$SSH_TEST_ROOT/activation-commands" should not be exist
         End
     End
 
-    It 'blocks a socket installed and enabled by a later package migration'
-        migrate() {
-            printf loaded > "$MODEL/ssh.service.LoadState"
-            printf enabled > "$MODEL/ssh.service.UnitFileState"
-            disableSSH
-            printf loaded > "$MODEL/ssh.socket.LoadState"
-            printf enabled > "$MODEL/ssh.socket.UnitFileState"
-            printf disabled > "$MODEL/ssh.service.UnitFileState"
-            systemctl start ssh.socket
-        }
-        When call migrate
+    It 'allows explicit legacy activation without consuming an unowned record'
+        mkdir -p "$SSH_POLICY_DIR"
+        printf 'sshd.service enabled\n' > "$SSH_POLICY_DIR/activation"
+        printf inactive > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+        When call restoreSSH
         The status should be success
-        The contents of file "$MODEL/ssh.socket.ActiveState" should equal inactive
-        The contents of file "$MODEL/ssh.socket.UnitFileState" should equal enabled
+        The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal active
+        The contents of file "$SSH_TEST_ROOT/activation-commands" should equal 'start sshd.service'
+        The contents of file "$SSH_POLICY_DIR/activation" should equal 'sshd.service enabled'
+    End
+
+    Describe 'socket backing services'
+        Parameters
+            ssh no ssh.service
+            sshd no custom-ssh.service
+            sshd yes sshd@internal.service
+        End
+
+        It "preserves the block when $1.socket uses a masked $3"
+            mask_socket_service() {
+                local socket="$1.socket" accept="$2" service="$3"
+                printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+                printf loaded > "$SSH_MODEL_DIR/$socket.LoadState"
+                printf enabled > "$SSH_MODEL_DIR/$socket.UnitFileState"
+                save_stopped_policy
+                printf '%s' "$accept" > "$SSH_MODEL_DIR/$socket.Accept"
+                printf '%s' "$service" > "$SSH_MODEL_DIR/$socket.Triggers"
+                printf masked > "$SSH_MODEL_DIR/$service.LoadState"
+                restoreSSH
+            }
+            When call mask_socket_service "$1" "$2" "$3"
+            The status should be failure
+            The stderr should include "$3 is masked"
+            The file "$SSH_POLICY_DIR/disabled" should be exist
+            The file "$SSH_POLICY_DIR/activation" should be exist
+            The file "$SSH_TEST_ROOT/activation-commands" should not be exist
+        End
     End
 
-    It 'accepts a successful stop of an already-failed SSH service'
-        stop_failed() {
-            printf loaded > "$MODEL/sshd.service.LoadState"
-            printf failed > "$MODEL/sshd.service.ActiveState"
-            disableSSH
+    It 'restores an accepting socket without starting a template instance'
+        restore_accepting_socket() {
+            printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            printf loaded > "$SSH_MODEL_DIR/sshd.socket.LoadState"
+            printf enabled > "$SSH_MODEL_DIR/sshd.socket.UnitFileState"
+            printf yes > "$SSH_MODEL_DIR/sshd.socket.Accept"
+            printf loaded > "$SSH_MODEL_DIR/sshd@internal.service.LoadState"
+            save_stopped_policy
+            restoreSSH
+            cat "$SSH_TEST_ROOT/activation-commands"
         }
-        When run stop_failed
+        When call restore_accepting_socket
         The status should be success
-        The file "$SSH_POLICY_DIR/disabled" should be exist
-        The contents of file "$MODEL/sshd.service.ActiveState" should equal failed
-        The contents of file "$TEST_ROOT/commands" should include 'stop sshd'
+        The output should equal "enable sshd.socket
+start sshd.socket"
+        The contents of file "$SSH_MODEL_DIR/sshd@internal.service.ActiveState" should equal inactive
     End
 
-    It 'stops an active masked service without removing its mask'
-        stop_masked() {
-            printf masked > "$MODEL/sshd.service.LoadState"
-            printf masked > "$MODEL/sshd.service.UnitFileState"
-            printf active > "$MODEL/sshd.service.ActiveState"
-            disableSSH
-        }
-        masked_stop_event() {
-            jq -se '[.[].Message | fromjson |
-                select(.Unit == "sshd" and .Phase == "Completed")] |
-                length == 1 and .[0].Skipped == false and .[0].StopExitCode == 0' \
-                "$EVENTS_LOGGING_DIR"/*.json
+    It 'does not guess a backing service when a non-accepting socket reports none'
+        restore_unknown_socket() {
+            printf disabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+            printf loaded > "$SSH_MODEL_DIR/sshd.socket.LoadState"
+            printf enabled > "$SSH_MODEL_DIR/sshd.socket.UnitFileState"
+            save_stopped_policy
+            : > "$SSH_MODEL_DIR/sshd.socket.Triggers"
+            restoreSSH
         }
-        When call stop_masked
-        The status should be success
-        The contents of file "$MODEL/sshd.service.ActiveState" should equal inactive
-        The contents of file "$MODEL/sshd.service.UnitFileState" should equal masked
-        The contents of file "$TEST_ROOT/commands" should not include 'unmask '
-        The result of function masked_stop_event should equal true
+        When call restore_unknown_socket
+        The status should be failure
+        The stderr should include "no backing service reported"
+        The file "$SSH_POLICY_DIR/disabled" should be exist
+        The file "$SSH_TEST_ROOT/activation-commands" should not be exist
     End
 
-    It 'stops listeners but fails rather than overwriting an administrator drop-in'
-        conflict() {
-            printf loaded > "$MODEL/sshd.service.LoadState"
-            printf active > "$MODEL/sshd.service.ActiveState"
-            mkdir -p "$SSH_SYSTEMD_UNIT_DIR/ssh.socket.d"
-            printf custom > "$SSH_SYSTEMD_UNIT_DIR/ssh.socket.d/90-aks-ssh-access.conf"
-            disableSSH
+    It 'reinstates its block when the saved activation cannot start'
+        fail_activation() {
+            save_stopped_policy
+            FAIL_COMMAND='start sshd.service'
+            restoreSSH
         }
-        When run conflict
-        The status should equal 172
-        The stderr should include 'Refusing to overwrite'
-        The contents of file "$MODEL/sshd.service.ActiveState" should equal inactive
-        The contents of file "$SSH_SYSTEMD_UNIT_DIR/ssh.socket.d/90-aks-ssh-access.conf" should equal custom
+        When call fail_activation
+        The status should be failure
+        The stderr should include "Failed to restore SSH activation"
+        The file "$SSH_POLICY_DIR/disabled" should be exist
+        The contents of file "$SSH_MODEL_DIR/sshd.service.ActiveState" should equal inactive
     End
 End
diff --git a/spec/support/ssh_policy.sh b/spec/support/ssh_policy.sh
new file mode 100644
index 00000000000..9870f2f17d5
--- /dev/null
+++ b/spec/support/ssh_policy.sh
@@ -0,0 +1,98 @@
+#!/bin/bash
+
+setup_ssh_policy() {
+    SSH_TEST_ROOT=$(mktemp -d)
+    SSH_POLICY_DIR="$SSH_TEST_ROOT/policy"
+    SSH_SYSTEMD_UNIT_DIR="$SSH_TEST_ROOT/systemd"
+    SSH_MODEL_DIR="$SSH_TEST_ROOT/model"
+    EVENTS_LOGGING_DIR="$SSH_TEST_ROOT/events"
+    mkdir -p "$SSH_MODEL_DIR" "$EVENTS_LOGGING_DIR"
+    SSH_MODEL_OS=azurelinux
+    FAIL_COMMAND=""
+    DISABLE_SSH=true
+    DISABLE_PUBKEY_AUTH=false
+    SHOULD_RESTORE_SSH=false
+    EXPECTED_PUBKEY=""
+    SSHD_CONFIG_FILE="$SSH_TEST_ROOT/sshd_config"
+    printf 'PubkeyAuthentication yes\nMatch User entra\n  PubkeyAuthentication yes\n' > "$SSHD_CONFIG_FILE"
+    local unit
+    for unit in ssh.socket sshd.socket ssh.service sshd.service ssh@internal.service sshd@internal.service; do
+        printf not-found > "$SSH_MODEL_DIR/$unit.LoadState"
+        printf disabled > "$SSH_MODEL_DIR/$unit.UnitFileState"
+        printf inactive > "$SSH_MODEL_DIR/$unit.ActiveState"
+        printf '%s' "$unit" > "$SSH_MODEL_DIR/$unit.Id"
+    done
+    printf loaded > "$SSH_MODEL_DIR/sshd.service.LoadState"
+    printf enabled > "$SSH_MODEL_DIR/sshd.service.UnitFileState"
+    printf active > "$SSH_MODEL_DIR/sshd.service.ActiveState"
+    for unit in ssh.socket sshd.socket; do
+        printf no > "$SSH_MODEL_DIR/$unit.Accept"
+        printf '%s.service' "${unit%.socket}" > "$SSH_MODEL_DIR/$unit.Triggers"
+    done
+}
+
+cleanup_ssh_policy() { rm -rf -- "$SSH_TEST_ROOT"; }
+
+ssh_model_unit() {
+    local unit="$1"
+    case "$unit" in ssh|sshd) unit="$unit.service" ;; esac
+    if [ "$SSH_MODEL_OS" = ubuntu ] && [ "$unit" = sshd.service ]; then
+        unit=ssh.service
+    fi
+    printf '%s' "$unit"
+}
+
+systemctl() {
+    local action="$1" unit property
+    printf '%s\n' "$*" >> "$SSH_TEST_ROOT/commands"
+    [ "$*" != "$FAIL_COMMAND" ] || return 1
+    case "$action" in
+        enable|start) printf '%s\n' "$*" >> "$SSH_TEST_ROOT/activation-commands" ;;
+    esac
+    [ "$action" != daemon-reload ] || return 0
+    if [ "$action" = is-active ]; then
+        unit=$(ssh_model_unit "$3")
+        [ "$(cat "$SSH_MODEL_DIR/$unit.ActiveState")" = active ]
+        return
+    fi
+    unit=$(ssh_model_unit "$2")
+    case "$action" in
+        show)
+            property="${3#--property=}"
+            cat "$SSH_MODEL_DIR/$unit.$property"
+            ;;
+        cat) [ "$(cat "$SSH_MODEL_DIR/$unit.LoadState")" = loaded ] ;;
+        stop) printf inactive > "$SSH_MODEL_DIR/$unit.ActiveState" ;;
+        disable) printf disabled > "$SSH_MODEL_DIR/$unit.UnitFileState" ;;
+        enable) printf enabled > "$SSH_MODEL_DIR/$unit.UnitFileState" ;;
+        start)
+            [ "$(cat "$SSH_MODEL_DIR/$unit.LoadState")" = loaded ] || return 1
+            if [ -f "$SSH_POLICY_DIR/disabled" ] &&
+                grep -qx "ConditionPathExists=!$SSH_POLICY_DIR/disabled" \
+                    "$SSH_SYSTEMD_UNIT_DIR/$unit.d/90-aks-ssh-access.conf"; then
+                return 0
+            fi
+            if [ -n "$EXPECTED_PUBKEY" ]; then
+                grep -qx "PubkeyAuthentication $EXPECTED_PUBKEY" "$SSHD_CONFIG_FILE" || return 1
+            fi
+            printf active > "$SSH_MODEL_DIR/$unit.ActiveState"
+            ;;
+        unmask) printf loaded > "$SSH_MODEL_DIR/$unit.LoadState" ;;
+        reload-or-restart) return 0 ;;
+        *) return 1 ;;
+    esac
+}
+
+timeout() { shift; "$@"; }
+sleep() { :; }
+sshd() {
+    if [ "$1" = -T ]; then
+        awk 'BEGIN { value="yes" } /^Match/ { exit } /^PubkeyAuthentication / { value=$2 }
+            END { print "pubkeyauthentication " value }' "$3"
+    fi
+}
+install() {
+    [ "$1" != -d ] || return 0
+    while [ "$#" -gt 2 ]; do shift; done
+    command cp "$1" "$2"
+}