Update 'prep-grid-vm.sh'

prep-grid-vm.sh

@@ -36,11 +36,33 @@ setup_tf_users() {
 
 configure_ssh() {
     echo "Configuring SSH"
-    wget https://docs.grid.tf/threefold_public/public/raw/branch/master/sshd_config
-    mv sshd_config /etc/ssh/
-    systemctl restart ssh
+    ubuntu_version=$(lsb_release -rs)
+
+    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
+        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
+        cat > /etc/systemd/system/ssh.socket.d/port.conf <<EOF
+[Socket]
+ListenStream=
+ListenStream=34022
+EOF
+
+        systemctl daemon-reload
+        systemctl restart ssh.socket
+        systemctl enable ssh.socket
+    else
+        echo "Unsupported Ubuntu version: $ubuntu_version"
+        return 1
+    fi
 }
 
+
 update_system
 setup_tf_users
 configure_ssh