diff --git a/du_setup.sh b/du_setup.sh
index 5716fc8..5669e00 100644
--- a/du_setup.sh
+++ b/du_setup.sh
@@ -1,8 +1,9 @@
 #!/bin/bash
 
 # Debian 12 and Ubuntu Server Hardening Interactive Script
-# Version: 0.53 | 2025-07-01
+# Version: 0.54 | 2025-07-02
 # Changelog:
+# - v0.54: Fix for rollback_ssh_changes() - more reliable on newer Ubuntu 
 # - v0.53: Fix for test_backup() - was failing if run as non root sudo user
 # - v0.52: Roll-back SSH config on failure to configure SSH port, confirmed SSH config support for Ubuntu 24.10
 # - v0.51: corrected repo links
@@ -92,7 +93,7 @@ print_header() {
     echo -e "${CYAN}╔═════════════════════════════════════════════════════════════════╗${NC}"
     echo -e "${CYAN}║                                                                 ║${NC}"
     echo -e "${CYAN}║       DEBIAN/UBUNTU SERVER SETUP AND HARDENING SCRIPT           ║${NC}"
-    echo -e "${CYAN}║                      v0.53 | 2025-07-01                         ║${NC}"
+    echo -e "${CYAN}║                      v0.54 | 2025-07-02                         ║${NC}"
     echo -e "${CYAN}║                                                                 ║${NC}"
     echo -e "${CYAN}╚═════════════════════════════════════════════════════════════════╝${NC}"
     echo
@@ -658,40 +659,178 @@ EOF
 
 rollback_ssh_changes() {
     print_info "Rolling back SSH configuration changes to port $PREVIOUS_SSH_PORT..."
-    rm -f /etc/systemd/system/ssh.service.d/override.conf
-    rm -f /etc/systemd/system/ssh.socket.d/override.conf
-    rm -f /etc/ssh/sshd_config.d/99-hardening.conf
-    if [[ -f "$SSHD_BACKUP_FILE" ]]; then
-        cp "$SSHD_BACKUP_FILE" /etc/ssh/sshd_config
-        print_info "Restored original sshd_config from $SSHD_BACKUP_FILE."
+
+    # Ensure SSH_SERVICE is set and valid
+    local SSH_SERVICE=${SSH_SERVICE:-"sshd.service"}
+    local USE_SOCKET=false
+    # Check if socket activation is used
+    if systemctl list-units --full -all --no-pager | grep -E "[[:space:]]ssh.socket[[:space:]]" >/dev/null 2>&1; then
+        USE_SOCKET=true
+        SSH_SERVICE="ssh.socket"
+        print_info "Detected SSH socket activation: using ssh.socket."
+        log "Rollback: Using ssh.socket for SSH service."
+    elif ! systemctl list-units --full -all --no-pager | grep -E "[[:space:]]$SSH_SERVICE[[:space:]]" >/dev/null 2>&1; then
+        SSH_SERVICE="ssh.service" # Fallback for Ubuntu
+        print_warning "SSH service $SSH_SERVICE not found, falling back to ssh.service."
+        log "Rollback warning: Using fallback SSH service ssh.service."
+        # Verify fallback service exists
+        if ! systemctl list-units --full -all --no-pager | grep -E "[[:space:]]ssh.service[[:space:]]" >/dev/null 2>&1; then
+            print_error "No valid SSH service (sshd.service or ssh.service) found."
+            log "Rollback failed: No valid SSH service detected."
+            print_info "Action: Verify SSH service with 'systemctl list-units --full -all | grep ssh' and manually configure /etc/ssh/sshd_config."
+            return 0
+        fi
+    fi
+
+    # Remove systemd overrides for both service and socket
+    local OVERRIDE_DIR="/etc/systemd/system/${SSH_SERVICE}.d"
+    local SOCKET_OVERRIDE_DIR="/etc/systemd/system/ssh.socket.d"
+    local SERVICE_OVERRIDE_DIR="/etc/systemd/system/ssh.service.d"
+    if ! rm -rf "$OVERRIDE_DIR" "$SOCKET_OVERRIDE_DIR" "$SERVICE_OVERRIDE_DIR" 2>/dev/null; then
+        print_warning "Failed to remove systemd overrides at $OVERRIDE_DIR, $SOCKET_OVERRIDE_DIR, or $SERVICE_OVERRIDE_DIR."
+        log "Rollback warning: Failed to remove systemd overrides."
     else
-        print_error "Backup file not found. Could not restore sshd_config."
-        return 1
+        log "Removed systemd overrides: $OVERRIDE_DIR, $SOCKET_OVERRIDE_DIR, $SERVICE_OVERRIDE_DIR"
     fi
 
-    print_info "Reloading systemd and restarting $SSH_SERVICE..."
-    systemctl daemon-reload
-    if ! systemctl restart "$SSH_SERVICE"; then
-        print_warning "Failed to restart $SSH_SERVICE after rollback. Attempting manual start..."
-        /usr/sbin/sshd || true
+    # Remove custom SSH configuration
+    if ! rm -f /etc/ssh/sshd_config.d/99-hardening.conf 2>/dev/null; then
+        print_warning "Failed to remove /etc/ssh/sshd_config.d/99-hardening.conf."
+        log "Rollback warning: Failed to remove /etc/ssh/sshd_config.d/99-hardening.conf."
+    else
+        log "Removed /etc/ssh/sshd_config.d/99-hardening.conf"
     fi
 
+    # Restore original sshd_config
+    if [[ -f "$SSHD_BACKUP_FILE" ]]; then
+        if ! cp "$SSHD_BACKUP_FILE" /etc/ssh/sshd_config 2>/dev/null; then
+            print_error "Failed to restore sshd_config from $SSHD_BACKUP_FILE."
+            log "Rollback failed: Cannot copy $SSHD_BACKUP_FILE to /etc/ssh/sshd_config."
+            print_info "Action: Manually restore with 'cp $SSHD_BACKUP_FILE /etc/ssh/sshd_config' and verify with 'sshd -t'."
+            return 0
+        fi
+        print_info "Restored original sshd_config from $SSHD_BACKUP_FILE."
+        log "Restored sshd_config from $SSHD_BACKUP_FILE."
+    else
+        print_error "Backup file not found at $SSHD_BACKUP_FILE."
+        log "Rollback failed: $SSHD_BACKUP_FILE not found."
+        print_info "Action: Manually configure /etc/ssh/sshd_config to use port $PREVIOUS_SSH_PORT and verify with 'sshd -t'."
+        return 0
+    fi
+
+    # Validate restored sshd_config
+    if ! /usr/sbin/sshd -t >/tmp/sshd_config_test.log 2>&1; then
+        print_error "Restored sshd_config is invalid. Check /tmp/sshd_config_test.log for details."
+        log "Rollback failed: Invalid sshd_config after restoration. See /tmp/sshd_config_test.log."
+        print_info "Action: Fix /etc/ssh/sshd_config manually and test with 'sshd -t', then restart with 'systemctl restart ssh.service'."
+        return 0
+    fi
+
+    # Reload systemd
+    print_info "Reloading systemd..."
+    if ! systemctl daemon-reload 2>/dev/null; then
+        print_warning "Failed to reload systemd. Continuing with restart attempt..."
+        log "Rollback warning: Failed to reload systemd."
+    fi
+
+    # Handle socket activation or direct service restart
+    if [[ "$USE_SOCKET" == true ]]; then
+        # Stop ssh.socket to avoid conflicts
+        if systemctl is-active --quiet ssh.socket; then
+            if ! systemctl stop ssh.socket 2>/tmp/ssh_socket_stop.log; then
+                print_warning "Failed to stop ssh.socket. May affect port binding."
+                log "Rollback warning: Failed to stop ssh.socket. See /tmp/ssh_socket_stop.log."
+            else
+                log "Stopped ssh.socket to ensure correct port binding."
+            fi
+        fi
+        # Restart ssh.service to ensure sshd starts
+        print_info "Restarting ssh.service..."
+        if ! systemctl restart ssh.service 2>/tmp/sshd_restart.log; then
+            print_warning "Failed to restart ssh.service. Attempting manual start..."
+            log "Rollback warning: Failed to restart ssh.service. See /tmp/sshd_restart.log."
+            # Ensure no other sshd processes are running
+            pkill -f "sshd:.*" 2>/dev/null || true
+            # Manual start in foreground to verify
+            timeout 5 /usr/sbin/sshd -D -f /etc/ssh/sshd_config >/tmp/sshd_manual_start.log 2>&1
+            local TIMEOUT_EXIT=$?
+            if [[ $TIMEOUT_EXIT -eq 0 || $TIMEOUT_EXIT -eq 124 ]]; then
+                log "Manual SSH start succeeded (exit code $TIMEOUT_EXIT)."
+                # Restart ssh.service to ensure systemd management
+                if ! systemctl restart ssh.service 2>/tmp/sshd_restart_manual.log; then
+                    print_error "Failed to restart ssh.service after manual start."
+                    log "Rollback failed: Failed to restart ssh.service after manual start. See /tmp/sshd_restart_manual.log."
+                    print_info "Action: Check service status with 'systemctl status ssh.service' and logs with 'journalctl -u ssh.service'."
+                    return 0
+                fi
+            else
+                print_error "Manual SSH start failed (exit code $TIMEOUT_EXIT). Check /tmp/sshd_manual_start.log."
+                log "Rollback failed: Manual SSH start failed (exit code $TIMEOUT_EXIT). See /tmp/sshd_manual_start.log."
+                print_info "Action: Check service status with 'systemctl status ssh.service' and logs with 'journalctl -u ssh.service'."
+                return 0
+            fi
+        fi
+        # Restart ssh.socket to re-enable socket activation
+        print_info "Restarting ssh.socket..."
+        if ! systemctl restart ssh.socket 2>/tmp/ssh_socket_restart.log; then
+            print_warning "Failed to restart ssh.socket. SSH service may still be running."
+            log "Rollback warning: Failed to restart ssh.socket. See /tmp/ssh_socket_restart.log."
+        else
+            log "Restarted ssh.socket for socket activation."
+        fi
+    else
+        # Direct service restart for non-socket systems
+        print_info "Restarting $SSH_SERVICE..."
+        if ! systemctl restart "$SSH_SERVICE" 2>/tmp/sshd_restart.log; then
+            print_warning "Failed to restart $SSH_SERVICE. Attempting manual start..."
+            log "Rollback warning: Failed to restart $SSH_SERVICE. See /tmp/sshd_restart.log."
+            # Ensure no other sshd processes are running
+            pkill -f "sshd:.*" 2>/dev/null || true
+            # Manual start in foreground to verify
+            timeout 5 /usr/sbin/sshd -D -f /etc/ssh/sshd_config >/tmp/sshd_manual_start.log 2>&1
+            local TIMEOUT_EXIT=$?
+            if [[ $TIMEOUT_EXIT -eq 0 || $TIMEOUT_EXIT -eq 124 ]]; then
+                log "Manual SSH start succeeded (exit code $TIMEOUT_EXIT)."
+                # Restart service to ensure systemd management
+                if ! systemctl restart "$SSH_SERVICE" 2>/tmp/sshd_restart_manual.log; then
+                    print_error "Failed to restart $SSH_SERVICE after manual start."
+                    log "Rollback failed: Failed to restart $SSH_SERVICE after manual start. See /tmp/sshd_restart_manual.log."
+                    print_info "Action: Check service status with 'systemctl status $SSH_SERVICE' and logs with 'journalctl -u $SSH_SERVICE'."
+                    return 0
+                fi
+            else
+                print_error "Manual SSH start failed (exit code $TIMEOUT_EXIT). Check /tmp/sshd_manual_start.log."
+                log "Rollback failed: Manual SSH start failed (exit code $TIMEOUT_EXIT). See /tmp/sshd_manual_start.log."
+                print_info "Action: Check service status with 'systemctl status $SSH_SERVICE' and logs with 'journalctl -u $SSH_SERVICE'."
+                return 0
+            fi
+        fi
+    fi
+
+    # Verify rollback with retries
     local rollback_verified=false
     print_info "Verifying SSH rollback to port $PREVIOUS_SSH_PORT..."
     for ((i=1; i<=10; i++)); do
-        if ss -tuln | grep -q ":$PREVIOUS_SSH_PORT"; then
-            print_success "Rollback successful. SSH is now listening on port $PREVIOUS_SSH_PORT."
+        if ss -tuln | grep -q ":$PREVIOUS_SSH_PORT "; then
             rollback_verified=true
             break
         fi
-        sleep 1
+        log "Rollback verification attempt $i/10: SSH not listening on port $PREVIOUS_SSH_PORT."
+        sleep 3
     done
 
-    if [[ $rollback_verified == false ]]; then
+    if [[ $rollback_verified == true ]]; then
+        print_success "Rollback successful. SSH is now listening on port $PREVIOUS_SSH_PORT."
+        log "Rollback successful: SSH listening on port $PREVIOUS_SSH_PORT."
+    else
         print_error "Rollback failed. SSH service is not listening on port $PREVIOUS_SSH_PORT."
-        print_info "Please check service status manually with 'systemctl status $SSH_SERVICE' and 'journalctl -u $SSH_SERVICE'."
-        return 1
+        log "Rollback failed: SSH not listening on port $PREVIOUS_SSH_PORT. See /tmp/sshd_config_test.log, /tmp/sshd_restart.log, /tmp/sshd_manual_start.log, /tmp/ssh_socket_restart.log."
+        print_info "Action: Check service status with 'systemctl status ssh.service' or 'systemctl status ssh.socket', logs with 'journalctl -u ssh.service' or 'journalctl -u ssh.socket', and test config with 'sshd -t'."
+        print_info "Manually verify port with 'ss -tuln | grep :$PREVIOUS_SSH_PORT'."
+        print_info "Try starting SSH with 'sudo systemctl start ssh.service'."
     fi
+
+    return 0
 }
 
 configure_firewall() {