WizardKit/scripts/wk/clone/image.py

109 lines
2.3 KiB
Python

"""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.")