wrolpi/controller/lib/disks.py

"""
Disk utilities for WROLPi Controller.

Provides disk detection, mounting, and WROLPi drive identification.
"""

import json
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from controller.lib.config import get_config, is_docker_mode

SUPPORTED_FILESYSTEMS = {"ext4", "btrfs", "vfat", "exfat"}


@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)

    # 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"}

    # Safety rule: Don't unmount /media/wrolpi if no drives configured
    if mount_point == "/media/wrolpi":
        config = get_config()
        configured_mounts = config.get("drives", {}).get("mounts", [])

        if not configured_mounts:
            return {
                "success": False,
                "error": "Cannot unmount /media/wrolpi: No drives are configured. "
                         "This prevents accidental data loss.",
            }

        # Check if this is the primary drive
        primary_mount = next(
            (m for m in configured_mounts if m.get("is_primary")),
            None
        )
        if primary_mount:
            return {
                "success": False,
                "error": "Cannot unmount /media/wrolpi: This is the primary WROLPi drive.",
            }

    # 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 []