1 files changed, 10 insertions, 5 deletions

diff --git a/src/parted-usb.sh b/src/parted-usb.sh
index 198323b..d6d0873 100755
--- a/src/parted-usb.sh
+++ b/src/parted-usb.sh
@@ -13,17 +13,22 @@ MiB()
     KiB $(KiB "$1")
 }
 
-get_rootfs_size()
+get_size()
 {
-    [ "$ROOTFS" ] || return
-    if [ -b "$ROOTFS" ]
+    [ "$1" ] || return
+    if [ -b "$1" ]
     then
-        rootfs_size=$(blockdev --getsize64 "$ROOTFS")
+        blockdev --getsize64 "$1"
     else
-        rootfs_size=$(stat -c%s "$ROOTFS")
+        stat -c'%s' "$1"
     fi
 }
 
+get_rootfs_size()
+{
+    rootfs_size=$(get_size "$ROOTFS")
+}
+
 initialize_target()
 {
     if [ -b "$target" ]