diff --git a/prep-grid-vm.sh b/prep-grid-vm.sh
index b3e1fe4..8c5adf5 100644
--- a/prep-grid-vm.sh
+++ b/prep-grid-vm.sh
@@ -36,14 +36,14 @@ setup_tf_users() {
 
 configure_ssh() {
     echo "Configuring SSH"
-    ubuntu_version=$(lsb_release -rs)
+    ubuntu_version=$(lsb_release -rs 2>/dev/null || . /etc/os-release && echo "$VERSION_ID")
 
-    if [[ "$ubuntu_version" == "22.04" ]]; then
+    if [ "$ubuntu_version" = "22.04" ]; then
         echo "Detected Ubuntu 22.04 — replacing sshd_config and restarting ssh service"
         wget -q https://docs.grid.tf/threefold_public/public/raw/branch/master/sshd_config -O /etc/ssh/sshd_config
         systemctl restart ssh
 
-    elif [[ "$ubuntu_version" == "24.04" ]]; then
+    elif [ "$ubuntu_version" = "24.04" ]; then
         echo "Detected Ubuntu 24.04 — updating ssh.socket for port 34022"
         wget -q https://docs.grid.tf/threefold_public/public/raw/branch/master/sshd_config -O /etc/ssh/sshd_config
         mkdir -p /etc/systemd/system/ssh.socket.d
@@ -52,17 +52,16 @@ configure_ssh() {
 ListenStream=
 ListenStream=34022
 EOF
-
         systemctl daemon-reload
         systemctl restart ssh.socket
         systemctl enable ssh.socket
+
     else
         echo "Unsupported Ubuntu version: $ubuntu_version"
-        return 1
+        exit 1
     fi
 }
 
-
 update_system
 setup_tf_users
 configure_ssh