dwipe 2.0.0__py3-none-any.whl → 2.0.2__py3-none-any.whl

dwipe/DeviceInfo.py

@@ -1,18 +1,27 @@
 """
 DeviceInfo class for device discovery and information management
 """
+# pylint: disable=invalid-name,broad-exception-caught
+# pylint: disable=line-too-long,too-many-locals,too-many-branches
+# pylint: disable=too-many-return-statements,too-many-nested-blocks
+# pylint: disable=too-many-statements
+
+import os
+import re
 import json
 import subprocess
 import time
 import datetime
 import curses
+import traceback
 from fnmatch import fnmatch
 from types import SimpleNamespace
+from console_window import Theme
+from dataclasses import asdict
 
 from .WipeJob import WipeJob
 from .Utils import Utils
-from console_window import Theme
-from .PersistentState import PersistentState
+from .DrivePreChecker import DrivePreChecker
 
 
 class DeviceInfo:
@@ -21,8 +30,8 @@ class DeviceInfo:
 
     def __init__(self, opts, persistent_state=None):
         self.opts = opts
-        self.DB = opts.debug
-        self.wids = None
+        self.checker = DrivePreChecker()
+        self.wids = SimpleNamespace(state=5, name=4, human=7, fstype=4, label=5)
         self.head_str = None
         self.partitions = None
         self.persistent_state = persistent_state
@@ -36,6 +45,7 @@ class DeviceInfo:
                                dflt=dflt,         # default run-time state
                                label='',       # blkid
                                fstype='',      # blkid
+                               type='',        # device type (disk, part)
                                model='',       # /sys/class/block/{name}/device/vendor|model
                                size_bytes=size_bytes,  # /sys/block/{name}/...
                                marker='',      #  persistent status
@@ -45,10 +55,86 @@ class DeviceInfo:
                                job=None,         # if zap running
                                uuid='',        # filesystem UUID or PARTUUID
                                serial='',      # disk serial number (for whole disks)
+                               port='',        # port (for whole disks)
+                               hw_caps={},      # hw_wipe capabilities (for whole disks)
+                               hw_nopes={},   # hw reasons cannot do hw wipe
                                )
 
+    def get_hw_capabilities(self, ns):
+        """
+        Populates and returns hardware wipe capabilities for a disk.
+        Returns cached data if already present.
+        """
+        # 1. Check if we already have cached results
+        if hasattr(ns, 'hw_caps') and (ns.hw_caps or ns.hw_nopes):
+            return ns.hw_caps, ns.hw_nopes
+
+        # Initialize defaults
+        ns.hw_caps, ns.hw_nopes = {}, {}
+
+        # Skip hardware checks if firmware wipes are disabled
+        if not getattr(self.opts, 'firmware_wipes', False):
+            return ns.hw_caps, ns.hw_nopes
+
+        # 4. Perform the actual Probe
+        dev_path = f"/dev/{ns.name}"
+        if ns.name.startswith('nv'):
+            result = self.checker.check_nvme_drive(dev_path)
+        elif ns.name.startswith('sd'):
+            result = self.checker.check_ata_drive(dev_path)
+        # 5. Store Results
+        ns.hw_caps, ns.hw_nopes = result.modes, result.issues
+        return ns.hw_caps, ns.hw_nopes
+
+    def _get_port_from_sysfs(self, device_name):
+        try:
+            sysfs_path = f'/sys/class/block/{device_name}'
+            if not os.path.exists(sysfs_path):
+                return ''
+
+            real_path = os.path.realpath(sysfs_path).lower()
+
+            # 1. USB - Format: USB:1-1.4
+            if '/usb' in real_path:
+                usb_match = re.search(r'/(\d+-\d+(?:\.\d+)*):', real_path)
+                if usb_match:
+                    return f"USB:{usb_match.group(1)}"
+
+            # 2. SATA - Format: SATA:1
+            elif '/ata' in real_path:
+                ata_match = re.search(r'ata(\d+)', real_path)
+                if ata_match:
+                    return f"SATA:{ata_match.group(1)}"
+
+            # 3. NVMe - Format: PCI:1b.0 (Stripped of 0000:00: noise)
+            elif '/nvme' in real_path:
+                # This regex ignores the 4-digit domain and the first 2-digit bus
+                pci_match = re.search(r'0000:[0-9a-f]{2}:([0-9a-f]{2}\.[0-9a-f])', real_path)
+                if pci_match:
+                    return f"PCI:{pci_match.group(1)}"
+                return "NVMe"
+
+            # 4. MMC/eMMC - Format: MMC:0 or PCI:1a.0 (if PCI-attached)
+            elif '/mmc' in real_path:
+                # Try to extract mmc host number
+                mmc_match = re.search(r'/mmc_host/mmc(\d+)', real_path)
+                if mmc_match:
+                    return f"MMC:{mmc_match.group(1)}"
+                # Fallback: try to get PCI address if available
+                pci_match = re.search(r'0000:[0-9a-f]{2}:([0-9a-f]{2}\.[0-9a-f])', real_path)
+                if pci_match:
+                    return f"PCI:{pci_match.group(1)}"
+                return "MMC"
+
+        except Exception as e:
+            # Log exception to file for debugging
+            with open('/tmp/dwipe_port_debug.log', 'a', encoding='utf-8') as f:
+                f.write(f"Exception in _get_port_from_sysfs({device_name}): {e}\n")
+                traceback.print_exc(file=f)
+        return ''
+
     @staticmethod
-    def get_device_vendor_model(device_name):
+    def _get_device_vendor_model(device_name):
         """Gets the vendor and model for a given device from the /sys/class/block directory.
         - Args: - device_name: The device name, such as 'sda', 'sdb', etc.
         - Returns: A string containing the vendor and model information.
@@ -66,8 +152,15 @@ class DeviceInfo:
         rv = f'{get_str(device_name, "model")}'
         return rv.strip()
 
-    def parse_lsblk(self, dflt, prev_nss=None):
-        """Parse ls_blk for all the goodies we need"""
+    def parse_lsblk(self, dflt, prev_nss=None, lsblk_output=None):
+        """Parse ls_blk for all the goodies we need
+
+        Args:
+            dflt: Default state for new devices
+            prev_nss: Previous device namespaces for merging
+            lsblk_output: Optional lsblk JSON output string. If provided, uses this
+                         instead of running lsblk command. Useful for background monitoring.
+        """
         def eat_one(device):
             entry = self._make_partition_namespace(0, '', '', dflt)
             entry.name = device.get('name', '')
@@ -135,45 +228,88 @@ class DeviceInfo:
                         verify_prefix = '✓ '
                     elif verify_status == 'fail':
                         verify_prefix = '✗ '
-                    entry.marker = f'{verify_prefix}{state} {pct}% {marker.mode} {dt.strftime("%Y/%m/%d %H:%M")}'
+
+                    # Add error suffix if job failed abnormally
+                    error_suffix = ''
+                    abort_reason = getattr(marker, 'abort_reason', None)
+                    if abort_reason:
+                        error_suffix = f' Err[{abort_reason}]'
+
+                    entry.marker = f'{verify_prefix}{state} {pct}% {marker.mode} {dt.strftime("%Y/%m/%d %H:%M")}{error_suffix}'
                     entry.state = state
                     entry.dflt = state  # Set dflt so merge logic knows this partition has a marker
 
             return entry
 
-        # Run the `lsblk` command and get its output in JSON format with additional columns
-        # Use timeout to prevent UI freeze if lsblk hangs on problematic devices
-        try:
-            result = subprocess.run(['lsblk', '-J', '--bytes', '-o',
-                                    'NAME,MAJ:MIN,FSTYPE,TYPE,LABEL,PARTLABEL,FSUSE%,SIZE,MOUNTPOINTS,UUID,PARTUUID,SERIAL'],
-                                   stdout=subprocess.PIPE, text=True, check=False, timeout=10.0)
-            parsed_data = json.loads(result.stdout)
-        except subprocess.TimeoutExpired:
-            # lsblk hung - return empty dict to use previous device state
-            return {}
+        # Get lsblk output - either from parameter or by running command
+        if lsblk_output:  # Non-empty string from background monitor
+            # Use provided output string
+            try:
+                parsed_data = json.loads(lsblk_output)
+            except (json.JSONDecodeError, Exception):
+                # Invalid JSON - return empty dict
+                return {}
+        else:
+            # Run the `lsblk` command and get its output in JSON format with additional columns
+            # Use timeout to prevent UI freeze if lsblk hangs on problematic devices
+            try:
+                result = subprocess.run(['lsblk', '-J', '--bytes', '-o',
+                                        'NAME,MAJ:MIN,FSTYPE,TYPE,LABEL,PARTLABEL,FSUSE%,SIZE,MOUNTPOINTS,UUID,PARTUUID,SERIAL'],
+                                       stdout=subprocess.PIPE, text=True, check=False, timeout=10.0)
+                parsed_data = json.loads(result.stdout)
+            except (subprocess.TimeoutExpired, json.JSONDecodeError, Exception):  # pylint: disable=broad-exception-caught
+                # lsblk hung, returned bad JSON, or other error - return empty dict
+                # assemble_partitions will detect this and preserve previous state
+                return {}
         entries = {}
 
         # Parse each block device and its properties
         for device in parsed_data['blockdevices']:
             parent = eat_one(device)
-            parent.fstype = self.get_device_vendor_model(parent.name)
             entries[parent.name] = parent
             for child in device.get('children', []):
                 entry = eat_one(child)
                 entries[entry.name] = entry
                 entry.parent = parent.name
                 parent.minors.append(entry.name)
-                if not parent.fstype:
-                    parent.fstype = 'DISK'
                 self.disk_majors.add(entry.major)
                 if entry.mounts:
                     entry.state = 'Mnt'
                     parent.state = 'Mnt'
 
-        if self.DB:
-            print('\n\nDB: --->>> after parse_lsblk:')
-            for entry in entries.values():
-                print(vars(entry))
+
+        # Final pass: Identify disks, assign ports, and handle superfloppies
+        final_entries = {}
+        for name, entry in entries.items():
+            final_entries[name] = entry
+
+            # Only process top-level physical disks
+            if entry.parent is None:
+                # Hardware Info Gathering
+                entry.model = self._get_device_vendor_model(entry.name)
+                entry.port = self._get_port_from_sysfs(entry.name)
+
+                # The Split (Superfloppy Case)
+                # If it has children, the children already hold the data.
+                # If it has NO children but HAS data, we create the '----' child.
+                if not entry.minors and (entry.fstype or entry.label or entry.mounts):
+                    v_key = f"{name}_data"
+                    v_child = self._make_partition_namespace(entry.major, name, entry.size_bytes, dflt)
+                    v_child.name = "----"
+                    v_child.fstype = entry.fstype
+                    v_child.label = entry.label
+                    v_child.mounts = entry.mounts
+                    v_child.parent = name
+
+                    final_entries[v_key] = v_child
+                    entry.minors.append(v_key)
+
+                # Clean the hardware row of data-specific strings
+                entry.fstype = entry.model if entry.model else 'DISK'
+                entry.label = ''
+                entry.mounts = []
+
+        entries = final_entries
 
         return entries
 
@@ -199,9 +335,9 @@ class DeviceInfo:
 
         if to == 'STOP' and not state_in(ns.state, job_states):
             return False
-        if to == 'Lock' and not state_in(ns.state, list(ready_states) + ['Mnt']):
+        if to == 'Blk' and not state_in(ns.state, list(ready_states) + ['Mnt']):
             return False
-        if to == 'Unlk' and ns.state != 'Lock':
+        if to == 'Unbl' and ns.state != 'Blk':
             return False
 
         if to and fnmatch(to, '*%'):
@@ -221,7 +357,7 @@ class DeviceInfo:
         # Here we set inferences that block starting jobs
         #  -- clearing these states will be done on the device refresh
         if parent and state_in(ns.state, inferred_states):
-            if parent.state != 'Lock':
+            if parent.state != 'Blk':
                 parent.state = ns.state
         if state_in(ns.state, job_states):
             if parent:
@@ -230,6 +366,14 @@ class DeviceInfo:
                 minor.state = 'Busy'
         return True
 
+    @staticmethod
+    def clear_inferred_states(nss):
+        """Clear all inferred states (Busy, Mnt) so they can be re-inferred"""
+        inferred_states = ('Busy', 'Mnt')
+        for ns in nss.values():
+            if ns.state in inferred_states:
+                ns.state = ns.dflt
+
     @staticmethod
     def set_all_states(nss):
         """Set every state per linkage inferences"""
@@ -237,21 +381,40 @@ class DeviceInfo:
             DeviceInfo.set_one_state(nss, ns)
 
     def get_disk_partitions(self, nss):
-        """Filter out virtual/pseudo devices (zram, loop, etc.) while keeping physical drives.
-        Keeps: sd*, nvme*, vd*, hd*, mmcblk* (physical and USB drives)
-        Excludes: zram*, loop*, dm-*, sr*, ram*
-        """
-        # Virtual/pseudo device prefixes to exclude
-        exclude_prefixes = ('zram', 'loop', 'dm-', 'sr', 'ram')
+        """Filter to only wipeable physical storage using positive criteria.
 
+        Keeps devices that:
+        - Are type 'disk' or 'part' (from lsblk)
+        - Are writable (not read-only)
+        - Are real block devices (not virtual)
+
+        This automatically excludes:
+        - Virtual devices (zram, loop, dm-*, etc.)
+        - Read-only devices (CD-ROMs, eMMC boot partitions)
+        - Special partitions (boot loaders)
+        """
         ok_nss = {}
         for name, ns in nss.items():
             # Must be disk or partition type
             if ns.type not in ('disk', 'part'):
                 continue
+            if ns.size_bytes <= 0: # not relevant to wiping
+                continue
+
+            # Must be writable (excludes CD-ROMs, eMMC boot partitions, etc.)
+            ro_path = f'/sys/class/block/{name}/ro'
+            try:
+                with open(ro_path, 'r', encoding='utf-8') as f:
+                    if f.read().strip() != '0':
+                        continue  # Skip read-only devices
+            except (FileNotFoundError, Exception):
+                # If we can't read ro flag, skip this device to be safe
+                continue
 
-            # Exclude virtual/pseudo devices by name prefix
-            if any(name.startswith(prefix) for prefix in exclude_prefixes):
+            # Exclude common virtual device prefixes as a safety net
+            # (most should already be filtered by ro check or missing sysfs)
+            virtual_prefixes = ('zram', 'loop', 'dm-', 'ram')
+            if any(name.startswith(prefix) for prefix in virtual_prefixes):
                 continue
 
             # Include this device
@@ -261,7 +424,7 @@ class DeviceInfo:
 
     def compute_field_widths(self, nss):
         """Compute field widths for display formatting"""
-        wids = self.wids = SimpleNamespace(state=5, name=4, human=7, fstype=4, label=5)
+        wids = self.wids
         for ns in nss.values():
             wids.state = max(wids.state, len(ns.state))
             wids.name = max(wids.name, len(ns.name) + 2)
@@ -270,9 +433,6 @@ class DeviceInfo:
             wids.label = max(wids.label, len(ns.label))
             wids.fstype = max(wids.fstype, len(ns.fstype))
         self.head_str = self.get_head_str()
-        if self.DB:
-            print('\n\nDB: --->>> after compute_field_widths():')
-            print(f'self.wids={vars(wids)}')
 
     def get_head_str(self):
         """Generate header string for device list"""
@@ -357,8 +517,8 @@ class DeviceInfo:
         # Check for newly inserted flag first (hot-swapped devices should always show orange)
         if getattr(ns, 'newly_inserted', False):
             # Newly inserted device - orange/bright
-            if ns.state in ('Mnt', 'Lock'):
-                # Dim the orange for mounted/locked devices
+            if ns.state in ('Mnt', 'Blk'):
+                # Dim the orange for mounted/blocked devices
                 attr = curses.color_pair(Theme.HOTSWAP) | curses.A_DIM
             else:
                 attr = curses.color_pair(Theme.HOTSWAP) | curses.A_BOLD
@@ -377,8 +537,8 @@ class DeviceInfo:
         elif ns.state == '^':
             # Newly inserted device (hot-swapped) - orange/bright
             attr = curses.color_pair(Theme.HOTSWAP) | curses.A_BOLD
-        elif ns.state in ('Mnt', 'Lock'):
-            # Dim mounted or locked devices
+        elif ns.state in ('Mnt', 'Blk'):
+            # Dim mounted or blocked devices
             attr = curses.A_DIM
 
         # Override with red/danger color if verify failed
@@ -405,6 +565,7 @@ class DeviceInfo:
             if new_ns:
                 if prev_ns.job:
                     new_ns.job = prev_ns.job
+                # Note: Do NOT preserve port - use fresh value from current scan
                 new_ns.dflt = prev_ns.dflt
                 # Preserve the "wiped this session" flag
                 if hasattr(prev_ns, 'wiped_this_session'):
@@ -443,12 +604,12 @@ class DeviceInfo:
                         new_ns.verify_failed_msg = prev_ns.verify_failed_msg
                         new_ns.mounts = [prev_ns.verify_failed_msg]
 
-                if prev_ns.state == 'Lock':
-                    new_ns.state = 'Lock'
+                if prev_ns.state == 'Blk':
+                    new_ns.state = 'Blk'
                 elif new_ns.state not in ('s', 'W'):
                     new_ns.state = new_ns.dflt
                     # Don't copy forward percentage states (like "v96%") - only persistent states
-                    if prev_ns.state not in ('s', 'W', 'Busy', 'Unlk') and not prev_ns.state.endswith('%'):
+                    if prev_ns.state not in ('s', 'W', 'Busy', 'Unbl') and not prev_ns.state.endswith('%'):
                         new_ns.state = prev_ns.state  # re-infer these
             elif prev_ns.job:
                 # unplugged device with job..
@@ -463,30 +624,101 @@ class DeviceInfo:
                 new_ns.newly_inserted = True  # Mark for orange color even if locked/mounted
         return nss
 
-    def assemble_partitions(self, prev_nss=None):
-        """Assemble and filter partitions for display"""
-        nss = self.parse_lsblk(dflt='^' if prev_nss else '-', prev_nss=prev_nss)
+    def assemble_partitions(self, prev_nss=None, lsblk_output=None):
+        """Assemble and filter partitions for display
+
+        Args:
+            prev_nss: Previous device namespaces for merging
+            lsblk_output: Optional lsblk JSON output string from LsblkMonitor
+        """
+        nss = self.parse_lsblk(dflt='^' if prev_nss else '-', prev_nss=prev_nss,
+                               lsblk_output=lsblk_output)
+
+        # If parse_lsblk failed (returned empty) and we have previous data, keep previous state
+        if not nss and prev_nss:
+            # lsblk scan failed or returned no devices - preserve previous state
+            # This prevents losing devices when lsblk temporarily fails
+            # But clear temporary status messages from completed jobs
+            for ns in prev_nss.values():
+                if not ns.job and ns.mounts:
+                    # Job finished - clear temporary status messages like "Verified: zeroed"
+                    ns.mounts = [m for m in ns.mounts if not m.startswith(('Verified:', 'Stopped'))]
+            return prev_nss  # Return early - don't reprocess
 
         nss = self.get_disk_partitions(nss)
 
         nss = self.merge_dev_infos(nss, prev_nss)
 
-        # Apply persistent locked states
+        # Apply persistent blocked states
         if self.persistent_state:
             for ns in nss.values():
                 # Update last_seen timestamp
                 self.persistent_state.update_device_seen(ns)
-                # Apply persistent lock state
+                # Apply persistent block state
                 if self.persistent_state.get_device_locked(ns):
-                    ns.state = 'Lock'
+                    ns.state = 'Blk'
 
+        # Clear inferred states so they can be re-computed based on current job status
+        self.clear_inferred_states(nss)
         self.set_all_states(nss)  # set inferred states
 
         self.compute_field_widths(nss)
-
-        if self.DB:
-            print('\n\nDB: --->>> after assemble_partitions():')
-            for name, ns in nss.items():
-                print(f'DB: {name}: {vars(ns)}')
-        self.partitions = nss
         return nss
+
+    @staticmethod
+    def dump(parts=None, title='after lsblk'):
+        """Print nicely formatted device information"""
+        if not parts:
+            return
+
+        print(f'\n{"=" * 80}')
+        print(f'{title}')
+        print(f'{"=" * 80}\n')
+
+        # Separate disks and partitions
+        disks = {name: part for name, part in parts.items() if part.type == 'disk'}
+        partitions = {name: part for name, part in parts.items() if part.type == 'part'}
+
+        # Print each disk with its partitions
+        for disk_name in sorted(disks.keys()):
+            disk = disks[disk_name]
+
+            # Disk header
+            print(f'┌─ {disk.name} ({disk.model or "Unknown Model"})')
+            print(f'│  Size: {Utils.human(disk.size_bytes)}  Serial: {disk.serial or "N/A"}  Port: {disk.port or "N/A"}')
+            print(f'│  State: {disk.state}  Marker: {disk.marker or "(none)"}')
+
+            # Hardware capabilities
+            if disk.hw_caps:
+                caps = ', '.join(disk.hw_caps.keys())
+                print(f'│  Hardware: {caps}')
+
+            # Find and print partitions for this disk
+            disk_parts = [(name, part) for name, part in partitions.items()
+                         if part.parent == disk.name]
+
+            if disk_parts:
+                for i, (part_name, part) in enumerate(sorted(disk_parts)):
+                    is_last = (i == len(disk_parts) - 1)
+                    branch = '└─' if is_last else '├─'
+
+                    # Partition info
+                    label_str = f'"{part.label}"' if part.label else '(no label)'
+                    fstype_str = part.fstype or '(no filesystem)'
+
+                    print(f'│  {branch} {part.name}: {Utils.human(part.size_bytes)}')
+                    print(f'│  {"  " if is_last else "│ "}  Label: {label_str}  Type: {fstype_str}')
+                    print(f'│  {"  " if is_last else "│ "}  State: {part.state}  UUID: {part.uuid or "N/A"}')
+
+                    if part.marker:
+                        print(f'│  {"  " if is_last else "│ "}  Marker: {part.marker}')
+
+                    if part.mounts:
+                        mounts_str = ', '.join(part.mounts)
+                        print(f'│  {"  " if is_last else "│ "}  Mounted: {mounts_str}')
+            else:
+                print(f'│  └─ (no partitions)')
+
+            print('│')
+
+        print(f'{"=" * 80}\n')