SaraFem/Devices.py

"""Removable media section for the sidebar (Features/Major:
"Show mounted stuff, e.g. USB sticks being plugged in or out").

Talks to udisks2 over D-Bus (dbus-next, already used by PortalService.py),
the same service other file managers use:

- Plugged-in devices show up whether or not they're mounted yet. The old
  version only listed folders under /run/media/<user>, so a stick appeared
  only once something else (e.g. PCManFM) had mounted it, and only on the
  next 3-second poll.
- Changes arrive as udisks2 signals, so the list updates straight away.
- Clicking an unmounted device mounts it (no password for removable media)
  and opens it; the eject button unmounts and powers it off.
- With AutoMountDevices on, a newly plugged-in device is mounted right away.

Falls back to polling /run/media/<user> if udisks2 isn't reachable.
Everything runs on the page's own event loop (Page.run_task), so calling
.update() from here is safe."""

import asyncio
import getpass
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional

import flet as ft
from dbus_next import BusType, Message, MessageType, Variant
from dbus_next.aio import MessageBus

from Core import AppConfig, NotificationStyle, Palette

UDisks = "org.freedesktop.UDisks2"
UDisksPath = "/org/freedesktop/UDisks2"
BlockInterface = f"{UDisks}.Block"
FilesystemInterface = f"{UDisks}.Filesystem"
DriveInterface = f"{UDisks}.Drive"

PollInterval = 3.0  # seconds, only for the no-udisks2 fallback
MountRoots = [Path("/run/media") / getpass.getuser(), Path("/media") / getpass.getuser()]


@dataclass
class Device:
    Key: str  # udisks2 block object path, or the mount path in fallback mode
    Label: str
    MountPoint: Optional[Path]
    DrivePath: str = ""
    CanPowerOff: bool = False


def _HumanSize(Bytes: int) -> str:
    Size = float(Bytes)
    for Unit in ("B", "KB", "MB", "GB", "TB"):
        if Size < 1000 or Unit == "TB":
            return f"{Size:.0f} {Unit}" if Unit == "B" else f"{Size:.1f} {Unit}"
        Size /= 1000
    return f"{Bytes} B"


def _Value(Properties: dict, Name: str, Default=None):
    Item = Properties.get(Name)
    return Item.value if isinstance(Item, Variant) else Default


def _UserMount(MountPoint: Path) -> bool:
    return any(Root in MountPoint.parents for Root in MountRoots)


def ListFallbackDevices() -> list[Device]:
    Devices: list[Device] = []
    for Root in MountRoots:
        if not Root.is_dir():
            continue
        try:
            Devices += [Device(str(P), P.name, P) for P in Root.iterdir() if P.is_dir()]
        except OSError:
            pass
    return sorted(Devices, key=lambda D: D.Label.lower())


class UDisksClient:
    def __init__(self, Bus: MessageBus):
        self.Bus = Bus

    @classmethod
    async def Connect(cls) -> "UDisksClient":
        return cls(await MessageBus(bus_type=BusType.SYSTEM).connect())

    async def _Call(self, ObjectPath: str, Interface: str, Member: str, Signature: str = "", Body: Optional[list] = None):
        Reply = await self.Bus.call(Message(
            destination=UDisks, path=ObjectPath, interface=Interface, member=Member,
            signature=Signature, body=Body or [],
        ))
        if Reply.message_type == MessageType.ERROR:
            raise OSError(Reply.body[0] if Reply.body else Reply.error_name)
        return Reply.body

    async def ListDevices(self, ShowInternal: bool) -> list[Device]:
        Objects = (await self._Call(UDisksPath, "org.freedesktop.DBus.ObjectManager", "GetManagedObjects"))[0]
        Devices = []
        for ObjectPath, Interfaces in Objects.items():
            if FilesystemInterface not in Interfaces or BlockInterface not in Interfaces:
                continue
            Block = Interfaces[BlockInterface]
            DrivePath = _Value(Block, "Drive", "/")
            if _Value(Block, "HintIgnore", False) or DrivePath == "/":
                continue  # hidden, or not backed by a real drive (loop devices etc.)
            MountPoints = [
                Path(bytes(Raw).rstrip(b"\0").decode(errors="replace"))
                for Raw in _Value(Interfaces[FilesystemInterface], "MountPoints", [])
            ]
            MountPoint = MountPoints[0] if MountPoints else None
            if _Value(Block, "HintSystem", False):
                # Internal disks: only unmounted ones or ones mounted like
                # removable media, never "/", "/home", "/boot" and friends.
                if not ShowInternal or (MountPoint is not None and not _UserMount(MountPoint)):
                    continue
            Drive = Objects.get(DrivePath, {}).get(DriveInterface, {})
            Label = _Value(Block, "IdLabel", "") or _Value(Block, "HintName", "")
            if not Label:
                Model = " ".join(Part for Part in (_Value(Drive, "Vendor", ""), _Value(Drive, "Model", "")) if Part)
                Label = f"{_HumanSize(_Value(Block, 'Size', 0))} {Model or 'Volume'}".strip()
            Devices.append(Device(
                Key=ObjectPath,
                Label=Label,
                MountPoint=MountPoint,
                DrivePath=DrivePath,
                CanPowerOff=bool(_Value(Drive, "CanPowerOff", False)),
            ))
        return sorted(Devices, key=lambda D: D.Label.lower())

    async def Watch(self, OnChange: Callable[[], None]) -> None:
        """Call OnChange for any udisks2 signal: devices added/removed
        (InterfacesAdded/Removed) and mounts changing (PropertiesChanged)."""
        await self.Bus.call(Message(
            destination="org.freedesktop.DBus", path="/org/freedesktop/DBus",
            interface="org.freedesktop.DBus", member="AddMatch", signature="s",
            body=[f"type='signal',sender='{UDisks}'"],
        ))

        def Handler(Received: Message):
            if Received.message_type == MessageType.SIGNAL and (Received.path or "").startswith(UDisksPath):
                OnChange()
            return None

        self.Bus.add_message_handler(Handler)

    async def Mount(self, ObjectPath: str) -> Path:
        return Path((await self._Call(ObjectPath, FilesystemInterface, "Mount", "a{sv}", [{}]))[0])

    async def Unmount(self, ObjectPath: str) -> None:
        await self._Call(ObjectPath, FilesystemInterface, "Unmount", "a{sv}", [{}])

    async def PowerOff(self, DrivePath: str) -> None:
        await self._Call(DrivePath, DriveInterface, "PowerOff", "a{sv}", [{}])


class DevicesPanel(ft.Container):
    def __init__(self, OnNavigate: Callable[[Path], None], Config: AppConfig):
        self.OnNavigate = OnNavigate
        self.Config = Config
        self.Client: Optional[UDisksClient] = None
        self.CurrentDevices: list[Device] = []
        self.CurrentPath: Optional[Path] = None
        self.DeviceTiles: dict[str, ft.ListTile] = {}
        self.RefreshScheduled = False
        self.Busy: set[str] = set()  # devices being mounted/ejected right now
        self.Unmounted = False

        self.List = ft.ListView(spacing=0)
        self.List.controls = self._BuildTiles()
        super().__init__(content=self.List, padding=0)

    def did_mount(self) -> None:
        self.Unmounted = False
        self.page.run_task(self._Start)

    def will_unmount(self) -> None:
        # The sidebar gets rebuilt when its settings change; close this panel's
        # D-Bus connection so old panels don't pile up listening to udisks2.
        self.Unmounted = True
        if self.Client is not None:
            self.Client.Bus.disconnect()
            self.Client = None

    async def _Start(self) -> None:
        try:
            self.Client = await UDisksClient.Connect()
            await self.Client.Watch(self._ScheduleRefresh)
        except Exception:
            self.Client = None
            await self._PollLoop()
            return
        await self._Refresh(IsStartup=True)

    async def _PollLoop(self) -> None:
        while not self.Unmounted:
            self._Show(ListFallbackDevices())
            await asyncio.sleep(PollInterval)

    def _ScheduleRefresh(self) -> None:
        # One plug-in fires a burst of signals; refresh once per burst.
        if not self.RefreshScheduled and not self.Unmounted:
            self.RefreshScheduled = True
            self.page.run_task(self._DelayedRefresh)

    async def _DelayedRefresh(self) -> None:
        await asyncio.sleep(0.25)
        self.RefreshScheduled = False
        if self.Client is not None and not self.Unmounted:
            await self._Refresh()

    async def _Refresh(self, IsStartup: bool = False) -> None:
        try:
            Devices = await self.Client.ListDevices(self.Config.ShowInternalDrives)
        except Exception:
            return
        Known = {D.Key for D in self.CurrentDevices}
        self._Show(Devices, Force=IsStartup)
        if IsStartup or not self.Config.AutoMountDevices:
            return
        for NewDevice in Devices:
            if NewDevice.Key not in Known and NewDevice.MountPoint is None:
                await self._MountDevice(NewDevice, Open=False)

    def _Show(self, Devices: list[Device], Force: bool = False) -> None:
        if Devices == self.CurrentDevices and not Force:
            return
        self.CurrentDevices = Devices
        self.List.controls = self._BuildTiles()
        if self.CurrentPath is not None:
            self._ApplySelection()
        self.update()

    def _BuildTiles(self) -> list[ft.Control]:
        self.DeviceTiles = {}
        if not self.CurrentDevices:
            return [ft.Container(ft.Text("No devices.", color=Palette.Overlay0, size=12), padding=8)]

        Tiles = []
        for Current in self.CurrentDevices:
            Mounted = Current.MountPoint is not None
            Tile = ft.ListTile(
                title=Current.Label,
                leading=ft.Icon(
                    ft.Icons.USB if Mounted else ft.Icons.USB_OUTLINED,
                    color=Palette.Mauve if Mounted else Palette.Overlay0,
                ),
                trailing=ft.IconButton(
                    icon=ft.Icons.EJECT,
                    icon_color=Palette.Overlay0,
                    tooltip="Eject",
                    visible=self.Client is not None and (Mounted or Current.CanPowerOff),
                    on_click=lambda e, D=Current: self.page.run_task(self._EjectDevice, D),
                ),
                tooltip=str(Current.MountPoint) if Mounted else "Not mounted: click to mount",
                on_click=lambda e, D=Current: self.page.run_task(self._OpenDevice, D),
                selected_tile_color=Palette.Surface1,
                selected_color=Palette.Mauve,
            )
            self.DeviceTiles[Current.Key] = Tile
            Tiles.append(Tile)
        return Tiles

    async def _OpenDevice(self, Target: Device) -> None:
        if Target.MountPoint is not None:
            self.OnNavigate(Target.MountPoint)
        else:
            await self._MountDevice(Target, Open=True)

    async def _MountDevice(self, Target: Device, Open: bool) -> None:
        if self.Client is None or Target.Key in self.Busy:
            return
        self.Busy.add(Target.Key)
        try:
            MountPoint = await self.Client.Mount(Target.Key)
        except OSError as Error:
            self._Notify(f"Couldn't mount {Target.Label}: {Error}")
            return
        finally:
            self.Busy.discard(Target.Key)
        await self._Refresh()
        if Open:
            self.OnNavigate(MountPoint)

    async def _EjectDevice(self, Target: Device) -> None:
        if self.Client is None or Target.Key in self.Busy:
            return
        self.Busy.add(Target.Key)
        try:
            if Target.MountPoint is not None:
                # Step out of the device first, or the unmount is "busy".
                if self.CurrentPath is not None and (
                    self.CurrentPath == Target.MountPoint or Target.MountPoint in self.CurrentPath.parents
                ):
                    self.OnNavigate(Path.home())
                await self.Client.Unmount(Target.Key)
            if Target.CanPowerOff:
                await self.Client.PowerOff(Target.DrivePath)
            self._Notify(f"{Target.Label} can be removed now.", Kind="success")
        except OSError as Error:
            self._Notify(f"Couldn't eject {Target.Label}: {Error}")
        finally:
            self.Busy.discard(Target.Key)
        await self._Refresh()

    def _Notify(self, Message: str, Kind: str = "error") -> None:
        Background, Accent = NotificationStyle(Kind)
        self.page.show_dialog(ft.SnackBar(
            content=ft.Text(Message, color=Accent),
            bgcolor=Background,
            shape=ft.RoundedRectangleBorder(radius=4, side=ft.BorderSide(1, Accent)),
            open=True,
        ))

    def UpdateSelection(self, CurrentPath: Path) -> None:
        self.CurrentPath = CurrentPath
        if self._ApplySelection():
            self.update()

    def _ApplySelection(self) -> bool:
        Changed = False
        for Current in self.CurrentDevices:
            Tile = self.DeviceTiles.get(Current.Key)
            if Tile is None:
                continue
            NewValue = Current.MountPoint is not None and (
                self.CurrentPath == Current.MountPoint or Current.MountPoint in self.CurrentPath.parents
            )
            if Tile.selected != NewValue:
                Tile.selected = NewValue
                Changed = True
        return Changed