From e7642bdc63a73e8dd3b2be2071529ddaf54eaa1f Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 13 Aug 2023 17:46:31 -0700 Subject: [PATCH] Move image sections to wk/clone/image.py --- scripts/wk/clone/__init__.py | 1 + scripts/wk/clone/image.py | 109 +++++++++++++++++++++++++++++++++++ scripts/wk/clone/state.py | 90 +---------------------------- 3 files changed, 111 insertions(+), 89 deletions(-) create mode 100644 scripts/wk/clone/image.py diff --git a/scripts/wk/clone/__init__.py b/scripts/wk/clone/__init__.py index 021a5479..d52254b3 100644 --- a/scripts/wk/clone/__init__.py +++ b/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 diff --git a/scripts/wk/clone/image.py b/scripts/wk/clone/image.py new file mode 100644 index 00000000..06b700a9 --- /dev/null +++ b/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.") diff --git a/scripts/wk/clone/state.py b/scripts/wk/clone/state.py index 3c229ab2..f5706935 100644 --- a/scripts/wk/clone/state.py +++ b/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.")