fix strip error on CommandOutput and fix desktop env detect. linux

custom_components/easy_computer_manager/computer/__init__.py

@@ -1,5 +1,4 @@
 import asyncio
-from enum import Enum
 from typing import Optional, Dict, Any
 
 import asyncssh
@@ -7,14 +6,9 @@ from asyncssh import SSHClientConnection
 from wakeonlan import send_magic_packet
 
 from custom_components.easy_computer_manager import const, LOGGER
-from custom_components.easy_computer_manager.computer_utils import parse_gnome_monitors_config, \
+from custom_components.easy_computer_manager.computer.common import OSType, CommandOutput
-    parse_pactl_audio_config, parse_bluetoothctl
+from custom_components.easy_computer_manager.computer.utils import parse_gnome_monitors_output, \
-
+    parse_pactl_output, parse_bluetoothctl
-
-class OSType(str, Enum):
-    WINDOWS = "windows"
-    LINUX = "linux"
-    MACOS = "macos"
 
 
 class Computer:
@@ -23,17 +17,18 @@ class Computer:
         """Initialize the Computer object."""
         self.host = host
         self.mac = mac
-        self._username = username
+        self.username = username
         self._password = password
-        self._port = port
+        self.port = port
-        self._dualboot = dualboot
+        self.dualboot = dualboot
 
-        self._operating_system: Optional[OSType] = None
+        self.operating_system: Optional[OSType] = None
-        self._operating_system_version: Optional[str] = None
+        self.operating_system_version: Optional[str] = None
-        self._windows_entry_grub: Optional[str] = None
+        self.desktop_environment: Optional[str] = None
-        self._monitors_config: Optional[Dict[str, Any]] = None
+        self.windows_entry_grub: Optional[str] = None
-        self._audio_config: Dict[str, Optional[Dict]] = {}
+        self.monitors_config: Optional[Dict[str, Any]] = None
-        self._bluetooth_devices: Dict[str, Any] = {}
+        self.audio_config: Dict[str, Optional[Dict]] = {}
+        self.bluetooth_devices: Dict[str, Any] = {}
 
         self._connection: Optional[SSHClientConnection] = None
 
@@ -44,9 +39,9 @@ class Computer:
         try:
             client = await asyncssh.connect(
                 self.host,
-                username=self._username,
+                username=self.username,
                 password=self._password,
-                port=self._port,
+                port=self.port,
                 known_hosts=None
             )
             asyncssh.set_log_level("ERROR")
@@ -65,43 +60,47 @@ class Computer:
         await self._renew_connection()
 
         async def update_operating_system():
-            self._operating_system = await self._detect_operating_system()
+            self.operating_system = await self._detect_operating_system()
 
         async def update_operating_system_version():
-            self._operating_system_version = (await self.run_action("operating_system_version")).get("output")
+            self.operating_system_version = (await self.run_action("operating_system_version")).output
+
+        async def update_desktop_environment():
+            self.desktop_environment = (await self.run_action("desktop_environment")).output
 
         async def update_windows_entry_grub():
-            self._windows_entry_grub = (await self.run_action("get_windows_entry_grub")).get("output")
+            self.windows_entry_grub = (await self.run_action("get_windows_entry_grub")).output
 
         async def update_monitors_config():
-            monitors_config = (await self.run_action("get_monitors_config")).get("output")
+            monitors_config = (await self.run_action("get_monitors_config")).output
-            if self._operating_system == OSType.LINUX:
+            if self.operating_system == OSType.LINUX:
                 # TODO: add compatibility for KDE/others
-                self._monitors_config = parse_gnome_monitors_config(monitors_config)
+                self.monitors_config = parse_gnome_monitors_output(monitors_config)
-            elif self._operating_system == OSType.WINDOWS:
+            elif self.operating_system == OSType.WINDOWS:
                 # TODO: implement for Windows
                 pass
 
         async def update_audio_config():
-            speakers_config = (await self.run_action("get_speakers")).get("output")
+            speakers_config = (await self.run_action("get_speakers")).output
-            microphones_config = (await self.run_action("get_microphones")).get("output")
+            microphones_config = (await self.run_action("get_microphones")).output
 
-            if self._operating_system == OSType.LINUX:
+            if self.operating_system == OSType.LINUX:
-                self._audio_config = parse_pactl_audio_config(speakers_config, microphones_config)
+                self.audio_config = parse_pactl_output(speakers_config, microphones_config)
-            elif self._operating_system == OSType.WINDOWS:
+            elif self.operating_system == OSType.WINDOWS:
                 # TODO: implement for Windows
                 pass
 
         async def update_bluetooth_devices():
             bluetooth_config = await self.run_action("get_bluetooth_devices")
-            if self._operating_system == OSType.LINUX:
+            if self.operating_system == OSType.LINUX:
-                self._bluetooth_devices = parse_bluetoothctl(bluetooth_config)
+                self.bluetooth_devices = parse_bluetoothctl(bluetooth_config)
-            elif self._operating_system == OSType.WINDOWS:
+            elif self.operating_system == OSType.WINDOWS:
                 # TODO: implement for Windows
                 pass
 
         await update_operating_system()
         await update_operating_system_version()
+        await update_desktop_environment()
         await update_windows_entry_grub()
         await update_monitors_config()
         await update_audio_config()
@@ -110,7 +109,7 @@ class Computer:
     async def _detect_operating_system(self) -> OSType:
         """Detect the operating system of the computer."""
         uname_result = await self.run_manually("uname")
-        return OSType.LINUX if uname_result.get("return_code") == 0 else OSType.WINDOWS
+        return OSType.LINUX if uname_result.successful() else OSType.WINDOWS
 
     async def is_on(self, timeout: int = 1) -> bool:
         """Check if the computer is on by pinging it."""
@@ -123,38 +122,6 @@ class Computer:
         await proc.communicate()
         return proc.returncode == 0
 
-    def get_operating_system(self) -> OSType:
-        """Get the operating system of the computer."""
-        return self._operating_system
-
-    def get_operating_system_version(self) -> str:
-        """Get the operating system version of the computer."""
-        return self._operating_system_version
-
-    def get_windows_entry_grub(self) -> str:
-        """Get the Windows entry in the GRUB configuration file."""
-        return self._windows_entry_grub
-
-    def get_monitors_config(self) -> Dict[str, Any]:
-        """Get the monitors configuration of the computer."""
-        return self._monitors_config
-
-    def get_speakers(self) -> dict:
-        """Get the speakers configuration of the computer."""
-        return self._audio_config.get("speakers")
-
-    def get_microphones(self) -> dict:
-        """Get the microphones configuration of the computer."""
-        return self._audio_config.get("microphones")
-
-    def get_bluetooth_devices(self, return_as_string: bool = False) -> Dict[str, Any]:
-        """Get the Bluetooth devices of the computer."""
-        return self._bluetooth_devices
-
-    def is_dualboot(self) -> bool:
-        """Check if the computer is dualboot."""
-        return self._dualboot
-
     async def start(self) -> None:
         """Start the computer."""
         send_magic_packet(self.mac)
@@ -203,13 +170,13 @@ class Computer:
         # Implementation needed
         pass
 
-    async def run_action(self, action: str, params: Dict[str, Any] = None, exit: bool = None) -> Dict[str, Any]:
+    async def run_action(self, action: str, params: Dict[str, Any] = None, exit: bool = None) -> CommandOutput:
         """Run a predefined command via SSH."""
         if params is None:
             params = {}
 
         if action not in const.ACTIONS:
-            return {"output": "", "error": "Action not found", "return_code": 1}
+            return CommandOutput("", 1, "", "Action not found")
 
         command_template = const.ACTIONS[action]
 
@@ -219,32 +186,32 @@ class Computer:
         if "exit" in command_template and exit is None:
             exit = command_template["exit"]
 
-        commands = command_template.get(self._operating_system)
+        commands = command_template.get(self.operating_system.lower())
         if not commands:
-            raise ValueError(f"Action not supported for OS: {self._operating_system}")
+            raise ValueError(f"Action not supported for OS: {self.operating_system}")
 
-        result = None
+        command_result = None
         for command in commands:
             for param, value in params.items():
                 command = command.replace(f"%{param}%", value)
 
-            result = await self.run_manually(command)
+            command_result = await self.run_manually(command)
-            if result['return_code'] == 0:
+            if command_result.successful():
                 LOGGER.debug(f"Command successful: {command}")
-                return result
+                return command_result
 
             LOGGER.debug(f"Command failed: {command}")
 
         if exit:
             raise ValueError(f"Failed to run action: {action}")
 
-        return result
+        return command_result
 
-    async def run_manually(self, command: str) -> Dict[str, Any]:
+    async def run_manually(self, command: str) -> CommandOutput:
         """Run a custom command manually via SSH."""
         if not self._connection:
             await self._connect()
 
         result = await self._connection.run(command)
 
-        return {"output": result.stdout, "error": result.stderr, "return_code": result.exit_status}
+        return CommandOutput(command, result.exit_status, result.stdout, result.stderr)

custom_components/easy_computer_manager/computer/common.py

@@ -0,0 +1,18 @@
+from enum import Enum
+
+
+class OSType(str, Enum):
+    WINDOWS = "Windows"
+    LINUX = "Linux"
+    MACOS = "MacOS"
+
+
+class CommandOutput:
+    def __init__(self, command: str, return_code: int, output: str, error: str) -> None:
+        self.command = command
+        self.return_code = return_code
+        self.output = output.strip()
+        self.error = error.strip()
+
+    def successful(self) -> bool:
+        return self.return_code == 0

custom_components/easy_computer_manager/computer/utils.py

@@ -1,39 +1,41 @@
 import re
 
 from custom_components.easy_computer_manager import LOGGER
+from custom_components.easy_computer_manager.computer.common import CommandOutput
 
 
-async def format_debug_informations(computer: 'Computer'): # importing Computer causes circular import (how to fix?)
+async def format_debug_information(computer: 'Computer'):  # importing Computer causes circular import (how to fix?)
     """Return debug information about the host system."""
 
     data = {
         'os': {
-            'name': computer.get_operating_system(),
+            'name': computer.operating_system,
-            'version': computer.get_operating_system_version(),
+            'version': computer.operating_system_version,
+            'desktop_environment': computer.desktop_environment
         },
         'connection':{
             'host': computer.host,
             'mac': computer.mac,
-            'username': computer._username,
+            'username': computer.username,
-            'port': computer._port,
+            'port': computer.port,
-            'dualboot': computer._dualboot,
+            'dualboot': computer.dualboot,
             'is_on': await computer.is_on()
         },
         'grub':{
-            'windows_entry': computer.get_windows_entry_grub()
+            'windows_entry': computer.windows_entry_grub
         },
         'audio':{
-            'speakers': computer.get_speakers(),
+            'speakers': computer.audio_config.get('speakers'),
-            'microphones': computer.get_microphones()
+            'microphones': computer.audio_config.get('microphones')
         },
-        'monitors': computer.get_monitors_config(),
+        'monitors': computer.monitors_config,
-        'bluetooth_devices': computer.get_bluetooth_devices()
+        'bluetooth_devices': computer.bluetooth_devices
     }
 
     return data
 
 
-def parse_gnome_monitors_config(config: str) -> list:
+def parse_gnome_monitors_output(config: str) -> list:
     """
     Parse the GNOME monitors configuration.
 
@@ -93,7 +95,7 @@ def parse_gnome_monitors_config(config: str) -> list:
     return monitors
 
 
-def parse_pactl_audio_config(config_speakers: str, config_microphones: str) -> dict[str, list]:
+def parse_pactl_output(config_speakers: str, config_microphones: str) -> dict[str, list]:
     """
     Parse the pactl audio configuration.
 
@@ -139,7 +141,7 @@ def parse_pactl_audio_config(config_speakers: str, config_microphones: str) -> d
     return config
 
 
-def parse_bluetoothctl(command: dict, connected_devices_only: bool = True,
+def parse_bluetoothctl(command: CommandOutput, connected_devices_only: bool = True,
                        return_as_string: bool = False) -> list | str:
     """Parse the bluetoothctl info command.
 
@@ -148,20 +150,17 @@ def parse_bluetoothctl(command: dict, connected_devices_only: bool = True,
     :param connected_devices_only:
         Only return connected devices.
         Will return all devices, connected or not, if False.
-    :param return_as_string:
-        Return the devices as a string (i.e. for device info in the UI).
 
-    :type command: dict
+    :type command: :class: CommandOutput
     :type connected_devices_only: bool
-    :type return_as_string: bool
 
     :returns: str | list
         The parsed bluetooth devices.
 
     """
 
-    if command.get('return_code') != 0:
+    if not command.successful():
-        if command.get("output").__contains__("Missing device address argument"):  # Means no devices are connected
+        if command.output.__contains__("Missing device address argument"):  # Means no devices are connected
             return "" if return_as_string else []
         else:
             LOGGER.warning(f"Cannot retrieve bluetooth devices, make sure bluetoothctl is installed")
@@ -170,7 +169,7 @@ def parse_bluetoothctl(command: dict, connected_devices_only: bool = True,
     devices = []
     current_device = None
 
-    for line in command.get("output").split('\n'):
+    for line in command.output.split('\n'):
         if line.startswith('Device'):
             if current_device is not None:
                 devices.append({
@@ -197,7 +196,4 @@ def parse_bluetoothctl(command: dict, connected_devices_only: bool = True,
     if connected_devices_only:
         devices = [device for device in devices if device["connected"] == "yes"]
 
-    if return_as_string:
-        devices = "; ".join([f"{device['name']} ({device['address']})" for device in devices])
-
     return devices

custom_components/easy_computer_manager/const.py

@@ -21,6 +21,10 @@ ACTIONS = {
         "windows": ['for /f "tokens=1 delims=|" %i in (\'wmic os get Name ^| findstr /B /C:"Microsoft"\') do @echo %i'],
         "linux": ["awk -F'=' '/^NAME=|^VERSION=/{gsub(/\"/, \"\", $2); printf $2\" \"}\' /etc/os-release && echo", "lsb_release -a | awk '/Description/ {print $2, $3, $4}'"]
     },
+    "desktop_environment": {
+        "linux": ["for session in $(ls /usr/bin/*session 2>/dev/null); do basename $session | sed 's/-session//'; done | grep -E 'gnome|kde|xfce|mate|lxde|cinnamon|budgie|unity' | head -n 1"],
+        "windows": ["echo Windows"]
+    },
     "shutdown": {
         "windows": ["shutdown /s /t 0", "wmic os where Primary=TRUE call Shutdown"],
         "linux": ["sudo shutdown -h now", "sudo init 0", "sudo systemctl poweroff"]

custom_components/easy_computer_manager/switch.py

@@ -24,8 +24,8 @@ from homeassistant.helpers.config_validation import make_entity_service_schema
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import computer_utils
+from .computer import OSType, Computer
-from .computer import Computer, OSType
+from .computer.utils import format_debug_information
 from .const import SERVICE_RESTART_TO_WINDOWS_FROM_LINUX, SERVICE_PUT_COMPUTER_TO_SLEEP, \
     SERVICE_START_COMPUTER_TO_WINDOWS, SERVICE_RESTART_COMPUTER, SERVICE_RESTART_TO_LINUX_FROM_WINDOWS, \
     SERVICE_CHANGE_MONITORS_CONFIG, SERVICE_STEAM_BIG_PICTURE, SERVICE_CHANGE_AUDIO_CONFIG, SERVICE_DEBUG_INFO, DOMAIN
@@ -134,7 +134,7 @@ class ComputerSwitch(SwitchEntity):
             name=self._attr_name,
             manufacturer="Generic",
             model="Computer",
-            sw_version=self.computer.get_operating_system_version(),
+            sw_version=self.computer.operating_system_version,
             connections={(dr.CONNECTION_NETWORK_MAC, self.computer.mac)},
         )
 
@@ -207,7 +207,7 @@ class ComputerSwitch(SwitchEntity):
 
     async def debug_info(self) -> ServiceResponse:
         """Prints debug info."""
-        return await computer_utils.format_debug_informations(self.computer)
+        return await format_debug_information(self.computer)
 
     async def async_update(self) -> None:
         """Ping the computer to see if it is online and update the state."""
@@ -218,10 +218,10 @@ class ComputerSwitch(SwitchEntity):
         # Update the state attributes and the connection only if the computer is on
         if self._state:
             self._attr_extra_state_attributes = {
-                "operating_system": self.computer.get_operating_system(),
+                "operating_system": self.computer.operating_system,
-                "operating_system_version": self.computer.get_operating_system_version(),
+                "operating_system_version": self.computer.operating_system_version,
                 "mac_address": self.computer.mac,
                 "ip_address": self.computer.host,
-                "connected_devices": self.computer.get_bluetooth_devices(return_as_string=True),
+                "connected_devices": self.computer.bluetooth_devices,
             }