wrolpi/controller/lib/disks.py

"""
Disk utilities for WROLPi Controller.

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

import grp
import json
import logging
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

logger = logging.getLogger(__name__)

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


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):
            is_partition = dev.get("type") == "part"
            is_whole_disk = dev.get("type") == "disk" and not dev.get("children")
            if (is_partition or is_whole_disk) 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=dev.get("model") or 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 filesystems without native POSIX permissions, set ownership to
    # the wrolpi user via uid/gid mount options. Covers FAT family and
    # both NTFS drivers (ntfs-3g FUSE and the in-kernel ntfs3).
    if fstype in ("exfat", "vfat", "ntfs", "ntfs3"):
        uid, gid = get_wrolpi_uid_gid()
        options = f"{options},uid={uid},gid={gid}"

    # NTFS volumes are frequently left "dirty" (Windows Fast Startup, or an
    # unclean unplug), which makes the driver refuse to mount them.  ntfsfix -d
    # clears the dirty flag and replays the journal; it is a no-op on clean
    # volumes.  Failure is non-fatal: the mount below reports the real error.
    if fstype in ("ntfs", "ntfs3"):
        try:
            fix = subprocess.run(
                ["ntfsfix", "-d", device],
                capture_output=True, text=True, timeout=60,
            )
            if fix.returncode != 0:
                logger.warning("ntfsfix %s failed (mount will still be attempted): %s",
                               device, (fix.stderr or fix.stdout).strip())
        except (subprocess.TimeoutExpired, OSError) as e:
            logger.warning("Could not run ntfsfix on %s: %s", device, e)

    # Build mount command
    cmd = ["mount"]
    if fstype:
        cmd.extend(["-t", fstype])
    if options:
        cmd.extend(["-o", options])
    cmd.extend([device, mount_point])

    logger.info("Mounting %s to %s (fstype=%s)", device, mount_point, fstype)
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        if result.returncode == 0:
            logger.info("Mounted %s to %s successfully", device, mount_point)
            return {"success": True, "mount_point": mount_point}
        else:
            logger.warning("Failed to mount %s to %s: %s", device, mount_point, result.stderr.strip())
            return {"success": False, "error": result.stderr.strip()}
    except subprocess.TimeoutExpired:
        logger.warning("Mount timed out for %s to %s", device, mount_point)
        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)

    logger.info("Unmounting %s (lazy=%s)", mount_point, lazy)
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        if result.returncode == 0:
            logger.info("Unmounted %s successfully", mount_point)
            return {"success": True, "mount_point": mount_point}
        else:
            logger.warning("Failed to unmount %s: %s", mount_point, result.stderr.strip())
            return {"success": False, "error": result.stderr.strip()}
    except subprocess.TimeoutExpired:
        logger.warning("Unmount timed out for %s", mount_point)
        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 check_shadowed_data(target: str) -> Optional[dict]:
    """Look for user data at `target` that would be hidden by a mount over it.

    Returns a dict with `size_bytes` and `entries` when non-config content is
    present, or None when the target is empty (or contains only `config/`).

    `config/` is whitelisted because it is created transiently by
    `generate_certificates.sh` before the primary drive is mounted; that is
    expected, not user data.

    Returns None for non-existent targets.
    """
    target_path = Path(target)
    if not target_path.exists() or not target_path.is_dir():
        return None

    try:
        entries = [e.name for e in target_path.iterdir() if e.name != "config"]
    except OSError:
        return None
    if not entries:
        return None

    size_bytes = 0
    for entry in entries:
        entry_path = target_path / entry
        try:
            # Skip symlinks: a symlink-to-directory at the mount target would
            # otherwise be walked, inflating size_bytes with content from
            # outside the would-be-shadowed area (e.g. /etc, /). The symlink
            # itself is still listed in `entries` so the user is warned about
            # it; it just contributes 0 bytes.
            if entry_path.is_symlink():
                continue
            if entry_path.is_dir():
                for path in entry_path.rglob("*"):
                    try:
                        if path.is_file() and not path.is_symlink():
                            size_bytes += path.stat().st_size
                    except OSError:
                        # File disappeared mid-walk or unreadable; skip.
                        continue
            elif entry_path.is_file():
                size_bytes += entry_path.stat().st_size
        except OSError:
            continue

    return {"size_bytes": size_bytes, "entries": sorted(entries)}


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