Move image sections to wk/clone/image.py

scripts/wk/clone/__init__.py

@@ -2,5 +2,6 @@
 
 from . import block_pair
 from . import ddrescue
+from . import image
 from . import menus
 from . import state

scripts/wk/clone/image.py

@@ -0,0 +1,109 @@
+"""WizardKit: ddrescue TUI - State"""
+# vim: sts=2 sw=2 ts=2
+
+import atexit
+import logging
+import pathlib
+import plistlib
+import re
+
+from wk import exe
+from wk.std import PLATFORM
+from wk.ui import cli
+
+
+# STATIC VARIABLES
+LOG = logging.getLogger(__name__)
+
+
+# Functions
+def mount_raw_image(path) -> pathlib.Path:
+  """Mount raw image using OS specific methods, returns pathlib.Path."""
+  loopback_path = None
+
+  if PLATFORM == 'Darwin':
+    loopback_path = mount_raw_image_macos(path)
+  elif PLATFORM == 'Linux':
+    loopback_path = mount_raw_image_linux(path)
+
+  # Check
+  if not loopback_path:
+    cli.print_error(f'Failed to mount image: {path}')
+
+  # Register unmount atexit
+  atexit.register(unmount_loopback_device, loopback_path)
+
+  # Done
+  return loopback_path
+
+
+def mount_raw_image_linux(path) -> pathlib.Path:
+  """Mount raw image using losetup, returns pathlib.Path."""
+  loopback_path = None
+
+  # Mount using losetup
+  cmd = [
+    'sudo',
+    'losetup',
+    '--find',
+    '--partscan',
+    '--show',
+    path,
+    ]
+  proc = exe.run_program(cmd, check=False)
+
+  # Check result
+  if proc.returncode == 0:
+    loopback_path = proc.stdout.strip()
+
+  # Done
+  return loopback_path
+
+
+def mount_raw_image_macos(path) -> pathlib.Path:
+  """Mount raw image using hdiutil, returns pathlib.Path."""
+  loopback_path = None
+  plist_data = {}
+
+  # Mount using hdiutil
+  # plistdata['system-entities'][{}...]
+  cmd = [
+    'hdiutil', 'attach',
+    '-imagekey', 'diskimage-class=CRawDiskImage',
+    '-nomount',
+    '-plist',
+    '-readonly',
+    path,
+    ]
+  proc = exe.run_program(cmd, check=False, encoding=None, errors=None)
+
+  # Check result
+  try:
+    plist_data = plistlib.loads(proc.stdout)
+  except plistlib.InvalidFileException:
+    return None
+  for dev in plist_data.get('system-entities', []):
+    dev_path = dev.get('dev-entry', '')
+    if re.match(r'^/dev/disk\d+$', dev_path):
+      loopback_path = dev_path
+
+  # Done
+  return loopback_path
+
+
+def unmount_loopback_device(path) -> None:
+  """Unmount loopback device using OS specific methods."""
+  cmd = []
+
+  # Build OS specific cmd
+  if PLATFORM == 'Darwin':
+    cmd = ['hdiutil', 'detach', path]
+  elif PLATFORM == 'Linux':
+    cmd = ['sudo', 'losetup', '--detach', path]
+
+  # Unmount loopback device
+  exe.run_program(cmd, check=False)
+
+
+if __name__ == '__main__':
+  print("This file is not meant to be called directly.")

scripts/wk/clone/state.py

@@ -1,7 +1,6 @@
 """WizardKit: ddrescue TUI - State"""
 # vim: sts=2 sw=2 ts=2
 
-import atexit
 import datetime
 import json
 import logging
@@ -22,6 +21,7 @@ import pytz
 from wk import cfg, debug, exe, io, log, net, std
 from wk.clone import menus
 from wk.clone.block_pair import BlockPair
+from wk.clone.image import mount_raw_image
 from wk.hw import disk as hw_disk
 from wk.hw.smart import (
   check_attributes,
@@ -1321,80 +1321,6 @@ def get_working_dir(mode, destination, force_local=False) -> pathlib.Path:
   return working_dir
 
 
-def mount_raw_image(path) -> pathlib.Path:
-  """Mount raw image using OS specific methods, returns pathlib.Path."""
-  loopback_path = None
-
-  if PLATFORM == 'Darwin':
-    loopback_path = mount_raw_image_macos(path)
-  elif PLATFORM == 'Linux':
-    loopback_path = mount_raw_image_linux(path)
-
-  # Check
-  if not loopback_path:
-    cli.print_error(f'Failed to mount image: {path}')
-
-  # Register unmount atexit
-  atexit.register(unmount_loopback_device, loopback_path)
-
-  # Done
-  return loopback_path
-
-
-def mount_raw_image_linux(path) -> pathlib.Path:
-  """Mount raw image using losetup, returns pathlib.Path."""
-  loopback_path = None
-
-  # Mount using losetup
-  cmd = [
-    'sudo',
-    'losetup',
-    '--find',
-    '--partscan',
-    '--show',
-    path,
-    ]
-  proc = exe.run_program(cmd, check=False)
-
-  # Check result
-  if proc.returncode == 0:
-    loopback_path = proc.stdout.strip()
-
-  # Done
-  return loopback_path
-
-
-def mount_raw_image_macos(path) -> pathlib.Path:
-  """Mount raw image using hdiutil, returns pathlib.Path."""
-  loopback_path = None
-  plist_data = {}
-
-  # Mount using hdiutil
-  # plistdata['system-entities'][{}...]
-  cmd = [
-    'hdiutil', 'attach',
-    '-imagekey', 'diskimage-class=CRawDiskImage',
-    '-nomount',
-    '-plist',
-    '-readonly',
-    path,
-    ]
-  proc = exe.run_program(cmd, check=False, encoding=None, errors=None)
-
-  # Check result
-  try:
-    plist_data = plistlib.loads(proc.stdout)
-  except plistlib.InvalidFileException:
-    return None
-  for dev in plist_data.get('system-entities', []):
-    dev_path = dev.get('dev-entry', '')
-    if re.match(r'^/dev/disk\d+$', dev_path):
-      loopback_path = dev_path
-
-  # Done
-  return loopback_path
-
-
 def select_disk_obj(label:str, disk_menu: cli.Menu, disk_path: str) -> hw_disk.Disk:
   """Get disk based on path or menu selection, returns Disk."""
   if not disk_path:
@@ -1453,19 +1379,5 @@ def set_mode(docopt_args) -> str:
   return mode
 
 
-def unmount_loopback_device(path) -> None:
-  """Unmount loopback device using OS specific methods."""
-  cmd = []
-
-  # Build OS specific cmd
-  if PLATFORM == 'Darwin':
-    cmd = ['hdiutil', 'detach', path]
-  elif PLATFORM == 'Linux':
-    cmd = ['sudo', 'losetup', '--detach', path]
-
-  # Unmount loopback device
-  exe.run_program(cmd, check=False)
-
-
 if __name__ == '__main__':
   print("This file is not meant to be called directly.")