From 859bc990e01dc920be06f9050d2b26dfd436bda0 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 1 May 2022 16:53:24 -0700
Subject: [PATCH] Fix destination size checks

Addresses issue #166
---
 scripts/wk/clone/ddrescue.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py
index bfd3af05..b46cadaa 100644
--- a/scripts/wk/clone/ddrescue.py
+++ b/scripts/wk/clone/ddrescue.py
@@ -275,11 +275,17 @@ class BlockPair():
 
   def safety_check(self):
     """Run safety check and abort if necessary."""
+    # TODO: Expand section to support non-Linux systems
     dest_size = -1
-    if self.destination.exists():
-      dest_obj = hw_disk.Disk(self.destination)
-      dest_size = dest_obj.size
-      del dest_obj
+    if self.destination.is_block_device():
+      cmd = [
+        'lsblk', '--bytes', '--json',
+        '--nodeps', '--noheadings', '--output=size',
+        self.destination,
+        ]
+      json_data = exe.get_json_from_command(cmd)
+      dest_size = json_data['blockdevices'][0]['size']
+      del json_data
 
     # Check destination size if cloning
     if not self.destination.is_file() and dest_size < self.size:
@@ -712,13 +718,6 @@ class State():
     # Set mode
     self.mode = set_mode(docopt_args)
 
-    # Image mode is broken..
-    # TODO: Fix image mode
-    #       Definitely for Linux, maybe for macOS
-    if self.mode == 'Image':
-      std.print_error("I'm sorry but image mode is currently broken...")
-      std.abort()
-
     # Select source
     self.source = get_object(docopt_args['<source>'])
     if not self.source: