From e1bef0b150b20efd6758d49dd7fa7156221cf699 Mon Sep 17 00:00:00 2001
From: buildplan <170122315+buildplan@users.noreply.github.com>
Date: Mon, 13 Oct 2025 16:39:04 +0100
Subject: [PATCH] Refactor SSH service checks and error handling

---
 du_setup.sh | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/du_setup.sh b/du_setup.sh
index 0af163f..fc097cb 100644
--- a/du_setup.sh
+++ b/du_setup.sh
@@ -750,6 +750,7 @@ cleanup_and_exit() {
             print_error "Rollback failed. SSH may not be accessible. Please check 'systemctl status $SSH_SERVICE' and 'journalctl -u $SSH_SERVICE'."
         fi
     fi
+    trap - ERR
     exit $exit_code
 }
 
@@ -762,6 +763,7 @@ configure_ssh() {
     # Ensure openssh-server is installed
     if ! dpkg -l openssh-server | grep -q ^ii; then
         print_error "openssh-server package is not installed."
+        trap - ERR
         return 1
     fi
 
@@ -775,6 +777,7 @@ configure_ssh() {
         SSH_SERVICE="sshd.service"
     else
         print_error "No SSH service or daemon detected."
+        trap - ERR
         return 1
     fi
     print_info "Using SSH service: $SSH_SERVICE"
@@ -799,6 +802,7 @@ configure_ssh() {
         # Verify the key was added
         if [[ ! -s "$AUTH_KEYS" ]]; then
             print_error "Failed to create authorized_keys file."
+            trap - ERR
             return 1
         fi
         chmod 600 "$AUTH_KEYS"; chown -R "$USERNAME:$USERNAME" "$SSH_DIR"
@@ -808,15 +812,16 @@ configure_ssh() {
 
     print_warning "SSH Key Authentication Required for Next Steps!"
     echo -e "${CYAN}Test SSH access from a SEPARATE terminal now:${NC}"
-    if [[ "$SERVER_IP_V4" != "unknown" ]]; then
+    if [[ -n "$SERVER_IP_V4" && "$SERVER_IP_V4" != "unknown" ]]; then
         echo -e "${CYAN}  Using IPv4: ssh -p $CURRENT_SSH_PORT $USERNAME@$SERVER_IP_V4${NC}"
     fi
-    if [[ "$SERVER_IP_V6" != "not available" ]]; then
+    if [[ -n "$SERVER_IP_V6" && "$SERVER_IP_V6" != "not available" ]]; then
         echo -e "${CYAN}  Using IPv6: ssh -p $CURRENT_SSH_PORT $USERNAME@$SERVER_IP_V6${NC}"
     fi
 
     if ! confirm "Can you successfully log in using your SSH key?"; then
         print_error "SSH key authentication is mandatory to proceed."
+        trap - ERR
         return 1
     fi
 
@@ -858,12 +863,13 @@ EOF
         print_info "This may be due to existing configuration files on the system."
         if ! confirm "Continue despite configuration warnings?"; then
             print_error "Aborting SSH configuration."
-            rm /etc/ssh/sshd_config.d/99-hardening.conf
-            rm /etc/issue.net
-            rm -rf /etc/systemd/system/ssh.socket.d
-            rm -rf /etc/systemd/system/ssh.service.d
-            rm -rf /etc/systemd/system/sshd.service.d
+            rm -f /etc/ssh/sshd_config.d/99-hardening.conf
+            rm -f /etc/issue.net
+            rm -f /etc/systemd/system/ssh.socket.d/override.conf
+            rm -f /etc/systemd/system/ssh.service.d/override.conf
+            rm -f /etc/systemd/system/sshd.service.d/override.conf
             systemctl daemon-reload
+            trap - ERR
             return 1
         fi
     fi
@@ -873,6 +879,7 @@ EOF
     sleep 5
     if ! ss -tuln | grep -q ":$SSH_PORT"; then
         print_error "SSH not listening on port $SSH_PORT after restart!"
+        trap - ERR
         return 1
     fi
     print_success "SSH service restarted on port $SSH_PORT."
@@ -882,16 +889,17 @@ EOF
     sleep 2
     if ssh -p "$SSH_PORT" -o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=5 root@localhost true 2>/dev/null; then
         print_error "Root SSH login is still possible! Check configuration."
+        trap - ERR
         return 1
     else
         print_success "Confirmed: Root SSH login is disabled."
     fi
 
     print_warning "CRITICAL: Test new SSH connection in a SEPARATE terminal NOW!"
-    if [[ "$SERVER_IP_V4" != "unknown" ]]; then
+    if [[ -n "$SERVER_IP_V4" && "$SERVER_IP_V4" != "unknown" ]]; then
         print_info "Use IPv4: ssh -p $SSH_PORT $USERNAME@$SERVER_IP_V4"
     fi
-    if [[ "$SERVER_IP_V6" != "not available" ]]; then
+    if [[ -n "$SERVER_IP_V6" && "$SERVER_IP_V6" != "not available" ]]; then
         print_info "Use IPv6: ssh -p $SSH_PORT $USERNAME@$SERVER_IP_V6"
     fi
 
@@ -915,6 +923,7 @@ EOF
                 else
                     print_success "Rollback successful. SSH restored on original port $PREVIOUS_SSH_PORT."
                 fi
+                trap - ERR
                 return 1
             fi
         fi