WizardKit/scripts/wk/clone/menus.py

"""WizardKit: ddrescue TUI - Menus"""
# vim: sts=2 sw=2 ts=2

import logging
import pathlib

from wk.cfg.ddrescue import DDRESCUE_SETTINGS
from wk.hw.disk import Disk, get_disks
from wk.std import GenericAbort, PLATFORM, bytes_to_string
from wk.ui import ansi, cli


# STATIC VARIABLES
LOG = logging.getLogger(__name__)
CLONE_SETTINGS = {
  'Source': None,
  'Destination': None,
  'Create Boot Partition': False,
  'First Run': True,
  'Needs Format': False,
  'Table Type': None,
  'Partition Mapping': [
    # (5, 1) ## Clone source partition #5 to destination partition #1
    ],
  }
if PLATFORM == 'Darwin':
  # TODO: Direct I/O needs more testing under macOS
  DDRESCUE_SETTINGS['Default']['--idirect'] = {'Selected': False, 'Hidden': True}
  DDRESCUE_SETTINGS['Default']['--odirect'] = {'Selected': False, 'Hidden': True}
MENU_ACTIONS = (
  'Start',
  f'Change settings {ansi.color_string("(experts only)", "YELLOW")}',
  f'Detect drives {ansi.color_string("(experts only)", "YELLOW")}',
  'Quit')
MENU_TOGGLES = {
  'Auto continue (if recovery % over threshold)': True,
  'Retry (mark non-rescued sectors "non-tried")': False,
  }
SETTING_PRESETS = (
  'Default',
  'Fast',
  'Safe',
  )


# Functions
def main() -> cli.Menu:
  """Main menu, returns wk.ui.cli.Menu."""
  menu = cli.Menu(title=ansi.color_string('ddrescue TUI: Main Menu', 'GREEN'))
  menu.separator = ' '

  # Add actions, options, etc
  for action in MENU_ACTIONS:
    if not (PLATFORM == 'Darwin' and 'Detect drives' in action):
      menu.add_action(action)
  for toggle, selected in MENU_TOGGLES.items():
    menu.add_toggle(toggle, {'Selected': selected})

  # Done
  return menu


def settings(mode: str, silent: bool = True) -> cli.Menu:
  """Settings menu, returns wk.ui.cli.Menu."""
  title_text = [
    ansi.color_string('ddrescue TUI: Expert Settings', 'GREEN'),
    ' ',
    ansi.color_string(
      ['These settings can cause', 'MAJOR DAMAGE', 'to drives'],
      ['YELLOW', 'RED', 'YELLOW'],
      ),
    'Please read the manual before making changes',
    ]
  menu = cli.Menu(title='\n'.join(title_text))
  menu.separator = ' '
  preset = 'Default'
  if not silent:
    # Ask which preset to use
    cli.print_standard(
      f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}'
      )
    preset = cli.choice('Please select a preset:', SETTING_PRESETS)

    # Fix selection
    for _p in SETTING_PRESETS:
      if _p.startswith(preset):
        preset = _p

  # Add default settings
  menu.add_action('Load Preset')
  menu.add_action('Main Menu')
  for name, details in DDRESCUE_SETTINGS['Default'].items():
    menu.add_option(name, details.copy())

  # Update settings using preset
  if preset != 'Default':
    for name, details in DDRESCUE_SETTINGS[preset].items():
      menu.options[name].update(details.copy())

  # Disable direct output when saving to an image
  if mode == 'Image':
    menu.options['--odirect']['Disabled'] = True
    menu.options['--odirect']['Selected'] = False

  # Done
  return menu


def disks() -> cli.Menu:
  """Disk menu, returns wk.ui.cli.Menu()."""
  cli.print_info('Scanning disks...')
  available_disks = get_disks()
  menu = cli.Menu('ddrescue TUI: Disk selection')
  menu.disabled_str = 'Already selected'
  menu.separator = ' '
  menu.add_action('Quit')
  for disk in available_disks:
    menu.add_option(
      name=(
        f'{str(disk.path):<12} '
        f'{disk.bus:<5} '
        f'{bytes_to_string(disk.size, decimals=1, use_binary=False):<8} '
        f'{disk.model} '
        f'{disk.serial}'
        ),
      details={'Object': disk},
      )

  # Done
  return menu


def select_disk(prompt_msg: str, menu: cli.Menu) -> Disk:
  """Select disk from provided Menu, returns Disk()."""
  menu.title = ansi.color_string(
    f'ddrescue TUI: {prompt_msg} Selection', 'GREEN',
    )

  # Get selection
  selection = menu.simple_select()
  if 'Quit' in selection:
    raise GenericAbort()

  # Disable selected disk's menu entry
  menu.options[selection[0]]['Disabled'] = True

  # Update details to include child devices
  selected_disk = selection[-1]['Object']
  selected_disk.update_details(skip_children=False)

  # Done
  return selected_disk


def select_disk_parts(prompt_msg, disk) -> list[Disk]:
  """Select disk parts from list, returns list of Disk()."""
  title = ansi.color_string('ddrescue TUI: Partition Selection', 'GREEN')
  title += f'\n\nDisk: {disk.path} {disk.description}'
  menu = cli.Menu(title)
  menu.separator = ' '
  menu.add_action('All')
  menu.add_action('None')
  menu.add_action('Proceed', {'Separator': True})
  menu.add_action('Quit')
  object_list = []

  def _select_parts(menu) -> None:
    """Loop over selection menu until at least one partition selected."""
    while True:
      selection = menu.advanced_select(
        f'Please select the parts to {prompt_msg.lower()}: ',
        )
      if 'All' in selection:
        for option in menu.options.values():
          option['Selected'] = True
      elif 'None' in selection:
        for option in menu.options.values():
          option['Selected'] = False
      elif 'Proceed' in selection:
        if any(option['Selected'] for option in menu.options.values()):
          # At least one partition/device selected/device selected
          break
      elif 'Quit' in selection:
        raise GenericAbort()

  # Bail early if running under macOS
  if PLATFORM == 'Darwin':
    return [disk]

  # Bail early if child device selected
  if disk.parent:
    return [disk]

  # Add parts
  whole_disk_str = f'{str(disk.path):<14} (Whole device)'
  for part in disk.children:
    fstype = part.get('fstype', '')
    fstype = str(fstype) if fstype else ''
    size = part["size"]
    name = (
      f'{str(part["path"]):<14} '
      f'{fstype.upper():<5} '
      f'({bytes_to_string(size, decimals=1, use_binary=True):>6})'
      )
    menu.add_option(name, details={'Selected': True, 'pathlib.Path': part['path']})

  # Add whole disk if necessary
  if not menu.options:
    menu.add_option(whole_disk_str, {'Selected': True, 'pathlib.Path': disk.path})
    menu.title += '\n\n'
    menu.title += ansi.color_string(' No partitions detected.', 'YELLOW')

  # Get selection
  _select_parts(menu)

  # Build list of Disk() object_list
  for option in menu.options.values():
    if option['Selected']:
      object_list.append(option['pathlib.Path'])

  # Check if whole disk selected
  if len(object_list) == len(disk.children):
    # NOTE: This is not true if the disk has no partitions
    msg = f'Preserve partition table and unused space in {prompt_msg.lower()}?'
    if cli.ask(msg):
      # Replace part list with whole disk obj
      object_list = [disk.path]

  # Convert object_list to Disk() objects
  cli.print_standard(' ')
  cli.print_info('Getting disk/partition details...')
  object_list = [Disk(path) for path in object_list]

  # Done
  return object_list


def select_path(prompt_msg) -> pathlib.Path:
  """Select path, returns pathlib.Path."""
  invalid = False
  menu = cli.Menu(
    title=ansi.color_string(f'ddrescue TUI: {prompt_msg} Path Selection', 'GREEN'),
    )
  menu.separator = ' '
  menu.add_action('Quit')
  menu.add_option('Current directory')
  menu.add_option('Enter manually')
  path = pathlib.Path.cwd()

  # Make selection
  selection = menu.simple_select()
  if 'Current directory' in selection:
    pass
  elif 'Enter manually' in selection:
    path = pathlib.Path(cli.input_text('Please enter path: '))
  elif 'Quit' in selection:
    raise GenericAbort()

  # Check
  try:
    path = path.resolve()
  except TypeError:
    invalid = True
  if invalid or not path.is_dir():
    cli.print_error(f'Invalid path: {path}')
    raise GenericAbort()

  # Done
  return path


if __name__ == '__main__':
  print("This file is not meant to be called directly.")