diff --git a/includes/rootfs-expand.ksi b/includes/rootfs-expand.ksi
index cf08239..ee0548c 100644
--- a/includes/rootfs-expand.ksi
+++ b/includes/rootfs-expand.ksi
@@ -1,35 +1,31 @@
 %post
-cat >/usr/local/bin/rootfs-expand << EOF
+cat << EOF > /usr/local/bin/rootfs-expand
+#!/usr/bin/env bash
+set -euf -o pipefail
 
-#!/bin/bash
-clear
-part=\$(mount |grep '^/dev.* / ' |awk '{print \$1}')
-if [ -z "\$part" ];then
-    echo "Error detecting rootfs"
-    exit -1
+ROOT_PART=\$(mount | grep ' on / ' | awk '{print \$1}')
+if [[ \${ROOT_PART} =~ "mmcblk" || \${ROOT_PART} =~ "nvme" ]]; then
+    DEV=\$(echo "\${ROOT_PART}" | rev | sed -r 's/^.{2}//' | rev)
+    INTERMEDIATE_DEV="\${DEV}p"
+elif [[ \${ROOT_PART} =~ "vd" || \${ROOT_PART} =~ "sd" ]]; then
+    DEV=\$(echo "\${ROOT_PART}" | rev | sed -r 's/^.{1}//' | rev)
+    INTERMEDIATE_DEV="\${DEV}"
+else
+    >&2 echo "\$0: device type unsupported"
+    exit 1
 fi
-dev=\$(echo \$part|sed 's/[0-9]*\$//g')
-devlen=\${#dev}
-num=\${part:\$devlen}
-if [[ "\$dev" =~ ^/dev/mmcblk[0-9]*p\$ ]];then
-    dev=\${dev:0:-1}
-fi
-if [ ! -x /usr/bin/growpart ];then
-    echo "Please install cloud-utils-growpart (sudo yum install cloud-utils-growpart)"
-    exit -2
-fi
-if [ ! -x /usr/sbin/resize2fs ];then
-    echo "Please install e2fsprogs (sudo yum install e2fsprogs)"
-    exit -3
-fi
-echo \$part \$dev \$num
+PART_NUM="\${ROOT_PART#"\$INTERMEDIATE_DEV"}"
+
+echo "\${DEV} \${PART_NUM} \${ROOT_PART}"
+
+echo "Extending partition \${PART_NUM} of \${DEV} to max size..."
+growpart "\${DEV}" "\${PART_NUM}"
+
+echo "Resising ext4 filesystem on \${ROOT_PART}..."
+resize2fs "\${ROOT_PART}"
 
-echo "Extending partition \$num to max size ...."
-growpart \$dev \$num
-echo "Resizing ext4 filesystem ..."
-resize2fs \$part
 echo "Done."
-df -h |grep \$part
+df -kh | head -n 1 && df -kh | grep "\${DEV}" | sort
 EOF
 
 chmod +x /usr/local/bin/rootfs-expand