diff --git a/meta/scripts/lib/wic/plugins/source/rootfs-u-boot.py b/meta/scripts/lib/wic/plugins/source/rootfs-u-boot.py
index 3ac0aa0..1e10648 100644
--- a/meta/scripts/lib/wic/plugins/source/rootfs-u-boot.py
+++ b/meta/scripts/lib/wic/plugins/source/rootfs-u-boot.py
@@ -37,7 +37,7 @@ class RootfsUBootPlugin(RootfsPlugin):
         # Prologue from RootfsPlugin.do_prepare_partition, retrieves the
         # rootfs directory
         if part.rootfs_dir is None:
-            if not 'ROOTFS_DIR' in krootfs_dir:
+            if 'ROOTFS_DIR' not in krootfs_dir:
                 raise WicError("Couldn't find --rootfs-dir, exiting")
 
             rootfs_dir = krootfs_dir['ROOTFS_DIR']
@@ -48,7 +48,8 @@ class RootfsUBootPlugin(RootfsPlugin):
                 rootfs_dir = part.rootfs_dir
             else:
                 raise WicError("Couldn't find --rootfs-dir=%s connection or "
-                               "it is not a valid path, exiting" % part.rootfs_dir)
+                               "it is not a valid path, exiting" %
+                               part.rootfs_dir)
         if os.path.isdir(rootfs_dir):
             real_rootfs_dir = rootfs_dir
         else:
@@ -59,6 +60,13 @@ class RootfsUBootPlugin(RootfsPlugin):
                                (rootfs_dir, image_rootfs_dir))
             real_rootfs_dir = image_rootfs_dir
 
+        root_dev = cr.rootdev
+        if not root_dev:
+            root_dev = source_params.get("root", None)
+            if not root_dev:
+                raise WicError("root not defined, exiting.")
+            root_dev = root_dev.replace(":", "=")
+
         u_boot_script = os.path.join(real_rootfs_dir,
                                      "etc/default/u-boot-script")
         if not os.path.exists(u_boot_script):
@@ -68,8 +76,8 @@ class RootfsUBootPlugin(RootfsPlugin):
         with open(u_boot_script, 'w') as cfg:
             cfg.write('# Generated by wic, rootfs-u-boot plugin\n')
             cfg.write('ROOT_PARTITION="%d"\n' % part.realnum)
-            cfg.write('KERNEL_ARGS="root=%s %s"\n' % \
-                (cr.rootdev, cr.ks.bootloader.append or ""))
+            cfg.write('KERNEL_ARGS="root=%s %s"\n' %
+                      (root_dev, cr.ks.bootloader.append or ""))
             no_initrd = source_params.get('no_initrd') or ''
             cfg.write('NO_INITRD="%s"\n' % no_initrd)
             overlays = source_params.get('overlays') or ''