""" Disk utilities for WROLPi Controller. Provides disk detection, mounting, and WROLPi drive identification. """ import grp import json import pwd import subprocess import tempfile from dataclasses import dataclass from pathlib import Path from typing import Optional from controller.lib.config import is_docker_mode SUPPORTED_FILESYSTEMS = {"ext4", "btrfs", "vfat", "exfat"} def get_wrolpi_uid_gid() -> tuple[int, int]: """ Get uid and gid for the wrolpi user. Returns: Tuple of (uid, gid). Falls back to (1001, 1001) if wrolpi user/group doesn't exist. """ try: uid = pwd.getpwnam("wrolpi").pw_uid gid = grp.getgrnam("wrolpi").gr_gid return uid, gid except KeyError: return 1001, 1001 # wrolpi user is typically created after the first user @dataclass class BlockDevice: """Represents a block device/partition.""" name: str path: str size: str fstype: Optional[str] mountpoint: Optional[str] label: Optional[str] uuid: Optional[str] model: Optional[str] is_wrolpi_drive: bool = False def get_block_devices() -> list[BlockDevice]: """ Get all block devices using lsblk. Returns: list of BlockDevice objects """ if is_docker_mode(): return [] cmd = ["lsblk", "-J", "-o", "NAME,PATH,SIZE,TYPE,FSTYPE,MOUNTPOINT,LABEL,UUID,MODEL"] try: result = subprocess.run(cmd, capture_output=True, text=True, timeout=10) if result.returncode != 0: return [] data = json.loads(result.stdout) devices = [] def process_device(dev: dict, parent_model: str = None): if dev.get("type") == "part" and dev.get("fstype") in SUPPORTED_FILESYSTEMS: devices.append(BlockDevice( name=dev["name"], path=dev.get("path", f"/dev/{dev['name']}"), size=dev.get("size", ""), fstype=dev.get("fstype"), mountpoint=dev.get("mountpoint"), label=dev.get("label"), uuid=dev.get("uuid"), model=parent_model, )) # Process children (partitions) for child in dev.get("children", []): process_device(child, dev.get("model")) for device in data.get("blockdevices", []): process_device(device) return devices except (subprocess.SubprocessError, json.JSONDecodeError): return [] def get_uuid(device: str) -> Optional[str]: """ Get UUID for a device using blkid. Args: device: Device path (e.g., /dev/sda1) Returns: UUID string or None """ try: cmd = ["blkid", "-s", "UUID", "-o", "value", device] result = subprocess.run(cmd, capture_output=True, text=True, timeout=5) if result.returncode == 0: return result.stdout.strip() except subprocess.SubprocessError: pass return None def validate_mount_point(mount_point: str) -> None: """ Validate that mount point is safe. Args: mount_point: Path to validate Raises: ValueError: If mount point is not allowed """ import os if not mount_point.startswith("/media"): raise ValueError( f"Mount point must be under /media, got: {mount_point}. " "Mounting to other locations is not allowed." ) # Prevent path traversal normalized = os.path.normpath(mount_point) if not normalized.startswith("/media"): raise ValueError(f"Invalid mount point after normalization: {normalized}") def mount_drive( device: str, mount_point: str, fstype: Optional[str] = None, options: str = "defaults", ) -> dict: """ Mount a drive. Args: device: Device path or UUID mount_point: Where to mount fstype: Filesystem type (auto-detected if not specified) options: Mount options Returns: dict with success status """ if is_docker_mode(): return {"success": False, "error": "Not available in Docker mode"} # Validate mount point try: validate_mount_point(mount_point) except ValueError as e: return {"success": False, "error": str(e)} # Create mount point if needed Path(mount_point).mkdir(parents=True, exist_ok=True) # For exfat/vfat, add uid/gid options to set ownership to wrolpi user # These filesystems don't support POSIX permissions natively if fstype in ("exfat", "vfat"): uid, gid = get_wrolpi_uid_gid() options = f"{options},uid={uid},gid={gid}" # Build mount command cmd = ["mount"] if fstype: cmd.extend(["-t", fstype]) if options: cmd.extend(["-o", options]) cmd.extend([device, mount_point]) try: result = subprocess.run(cmd, capture_output=True, text=True, timeout=30) if result.returncode == 0: return {"success": True, "mount_point": mount_point} else: return {"success": False, "error": result.stderr.strip()} except subprocess.TimeoutExpired: return {"success": False, "error": "Mount timed out"} def unmount_drive(mount_point: str, lazy: bool = False) -> dict: """ Unmount a drive with safety checks. Args: mount_point: Path to unmount lazy: Use lazy unmount if busy Returns: dict with success status """ if is_docker_mode(): return {"success": False, "error": "Not available in Docker mode"} # Check if mount is busy if is_mount_busy(mount_point) and not lazy: return { "success": False, "error": f"Mount {mount_point} is busy. Use lazy unmount or stop services first.", } # Unmount cmd = ["umount"] if lazy: cmd.append("-l") cmd.append(mount_point) try: result = subprocess.run(cmd, capture_output=True, text=True, timeout=30) if result.returncode == 0: return {"success": True, "mount_point": mount_point} else: return {"success": False, "error": result.stderr.strip()} except subprocess.TimeoutExpired: return {"success": False, "error": "Unmount timed out"} def is_mount_busy(mount_point: str) -> bool: """Check if any processes are using the mount point.""" try: result = subprocess.run( ["lsof", "+D", mount_point], capture_output=True, text=True, timeout=10, ) return result.returncode == 0 and result.stdout.strip() except (subprocess.SubprocessError, FileNotFoundError): return False def get_mounts() -> list[dict]: """ Get current mounts under /media. Returns: list of mount info dicts """ try: result = subprocess.run( ["findmnt", "-J", "-l", "-o", "TARGET,SOURCE,FSTYPE,OPTIONS"], capture_output=True, text=True, timeout=10, ) if result.returncode != 0: return [] data = json.loads(result.stdout) mounts = [] for fs in data.get("filesystems", []): target = fs.get("target", "") if target.startswith("/media"): mounts.append({ "mount_point": target, "device": fs.get("source"), "fstype": fs.get("fstype"), "options": fs.get("options"), }) return mounts except (subprocess.SubprocessError, json.JSONDecodeError): return []