diff --git a/setup_harden_debian_ubuntu.sh b/setup_harden_debian_ubuntu.sh
index 123c612..ec11f73 100644
--- a/setup_harden_debian_ubuntu.sh
+++ b/setup_harden_debian_ubuntu.sh
@@ -577,17 +577,31 @@ configure_ssh() {
     SSHD_BACKUP_FILE="$BACKUP_DIR/sshd_config.backup_$(date +%Y%m%d_%H%M%S)"
     cp /etc/ssh/sshd_config "$SSHD_BACKUP_FILE"
 
-    # Use systemd drop-in for port override, preserving Ubuntu's socket activation
-    NEW_SSH_CONFIG=$(mktemp)
-    tee "$NEW_SSH_CONFIG" > /dev/null <<EOF
+    # Apply port override based on SSH service type
+    if [[ "$SSH_SERVICE" == "ssh.socket" ]]; then
+        print_info "Configuring SSH socket to listen on port $SSH_PORT..."
+        NEW_SSH_CONFIG=$(mktemp)
+        tee "$NEW_SSH_CONFIG" > /dev/null <<EOF
+[Socket]
+ListenStream=
+ListenStream=$SSH_PORT
+EOF
+        mkdir -p /etc/systemd/system/ssh.socket.d
+        mv "$NEW_SSH_CONFIG" /etc/systemd/system/ssh.socket.d/override.conf
+        chmod 644 /etc/systemd/system/ssh.socket.d/override.conf
+    else
+        print_info "Configuring SSH service to listen on port $SSH_PORT..."
+        NEW_SSH_CONFIG=$(mktemp)
+        tee "$NEW_SSH_CONFIG" > /dev/null <<EOF
 [Service]
 ExecStart=
 ExecStart=/usr/sbin/sshd -D -p $SSH_PORT
 EOF
-    mkdir -p /etc/systemd/system/ssh.service.d
-    mv "$NEW_SSH_CONFIG" /etc/systemd/system/ssh.service.d/override.conf
-    chmod 644 /etc/systemd/system/ssh.service.d/override.conf
-
+        mkdir -p /etc/systemd/system/ssh.service.d
+        mv "$NEW_SSH_CONFIG" /etc/systemd/system/ssh.service.d/override.conf
+        chmod 644 /etc/systemd/system/ssh.service.d/override.conf
+    fi
+    
     # Apply additional hardening via sshd_config.d
     NEW_SSH_CONFIG=$(mktemp)
     tee "$NEW_SSH_CONFIG" > /dev/null <<EOF