"""WizardKit: Volume functions""" # vim: sts=2 sw=2 ts=2 import logging import re from wk import os as wk_os from wk.cfg.hw import ( VOLUME_FAILURE_THRESHOLD, VOLUME_WARNING_THRESHOLD, VOLUME_SIZE_THRESHOLD, ) from wk.std import PLATFORM, bytes_to_string from wk.ui.ansi import color_string # STATIC VARIABLES LOG = logging.getLogger(__name__) # Functions def add_dev_line(test_obj, details) -> None: """Add device line to test report.""" if details is test_obj.dev: details = details.raw_details filesystem = details['fstype'] report_line = re.sub(r"^/dev/", " ", details['name']) label = details['label'] if label: label = f', "{color_string(label, "CYAN")}"' report_line += f' ({filesystem}{label if label else ""})' # Bail early if not filesystem: # Skip devices without a filesystem return # Get sizes used = -1 percent_used = -1 size = details['size'] if PLATFORM == 'Darwin': size = int(details.get('TotalSize', -1)) free = int(details.get('FreeSpace', 0)) used = size - free elif PLATFORM == 'Linux': free = details.get('fsavail', 0) used = details.get('fsused', -1) if free is None: free = 0 if used is None: used = -1 percent_used = (used / size) * 100 # Report Bitlocker if filesystem == 'BitLocker': test_obj.report.append(f'{report_line} {bytes_to_string(size)}') # Handle unsupported devices if not details['mountpoint']: # Under Linux the volume needs to be mounted to get used space used = -1 # Check for failures if (used > 0 and percent_used >= VOLUME_FAILURE_THRESHOLD and size >= VOLUME_SIZE_THRESHOLD * 1024**3): test_obj.failed = True # Build and color size_line if needed color = None if test_obj.failed: color = 'RED' elif percent_used >= VOLUME_WARNING_THRESHOLD: color = 'YELLOW' size_line = f'{bytes_to_string(size)}' if used > 0: size_line += f' ({bytes_to_string(used)} used, {percent_used:0.0f}% full)' size_line = color_string(size_line, str(color)) # Done test_obj.report.append(f'{report_line} {size_line}') def check_volume_utilization(test_obj) -> None: """Check volume utilization using OS specific methods.""" dev = test_obj.dev # Mount all volumes (read only if possible) mount_all_volumes(test_obj.dev) dev.update_details(skip_children=False) # Build report test_obj.report.append(color_string('Disk Utilization', 'BLUE')) for _d in (dev, *dev.children): add_dev_line(test_obj, _d) # Update test object if test_obj.failed: test_obj.dev.add_note('Full volume(s) detected', color='YELLOW') test_obj.passed = False test_obj.set_status('Failed') else: test_obj.passed = True test_obj.set_status('Passed') def mount_all_volumes(dev) -> None: """Mount all volumes for dev using.""" if PLATFORM == 'Darwin': # NOTE: Disabled due to a bug they should already be mounted by macOS #wk_os.mac.mount_disk(device_path=dev.path) pass elif PLATFORM == 'Linux': wk_os.linux.mount_volumes( device_path=dev.path, read_write=False, scan_corestorage=not any(t.failed for t in dev.tests), ) if __name__ == '__main__': print("This file is not meant to be called directly.")