Code style refactoring

custom_components/ssh_command/__init__.py

@@ -1,63 +1,73 @@
 import logging
 
 import voluptuous as vol
-
-from homeassistant.core import ServiceCall
 from homeassistant.const import CONF_HOST, CONF_PORT, CONF_USERNAME, CONF_PASSWORD
+from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import config_validation as cv
 from paramiko import AutoAddPolicy, RSAKey, SSHClient
 
 _LOGGER = logging.getLogger(__name__)
 
-DOMAIN = 'ssh_command'
+DOMAIN = "ssh_command"
 
 SSH_COMMAND_SCHEMA = vol.All(
-    vol.Schema({
-        vol.Optional(CONF_HOST): cv.string,
-        vol.Optional(CONF_PORT): cv.string,
-        vol.Optional(CONF_USERNAME): cv.string,
-        vol.Optional(CONF_PASSWORD): cv.string,
-    }, extra=vol.PREVENT_EXTRA))
-
-CONFIG_SCHEMA = vol.Schema({
-    DOMAIN: SSH_COMMAND_SCHEMA
-}, extra=vol.ALLOW_EXTRA)
-
-def setup(hass, hass_config):
-    hass.data[DOMAIN] = hass_config.get(DOMAIN, {})
-    async def exec_command(call: ServiceCall):
-        host = call.data.get('host', hass.data[DOMAIN].get(CONF_HOST, '172.17.0.1'))
-        port = call.data.get('port', hass.data[DOMAIN].get(CONF_PORT, 22))
-        username = call.data.get('user', hass.data[DOMAIN].get(CONF_USERNAME, 'pi'))
-        password = call.data.get('pass', hass.data[DOMAIN].get(CONF_PASSWORD, 'raspberry'))
-        command = call.data.get('command')
-        ssh_private_key_path = call.data.get('ssh_private_key_path')
+    vol.Schema(
+        {
+            vol.Optional(CONF_HOST): cv.string,
+            vol.Optional(CONF_PORT): cv.string,
+            vol.Optional(CONF_USERNAME): cv.string,
+            vol.Optional(CONF_PASSWORD): cv.string,
+            vol.Optional("private_key"): cv.string,
+        },
+        extra=vol.PREVENT_EXTRA,
+    )
+)
+
+CONFIG_SCHEMA = vol.Schema({DOMAIN: SSH_COMMAND_SCHEMA}, extra=vol.ALLOW_EXTRA)
+
+
+async def async_setup(hass: HomeAssistant, config: dict) -> bool:
+    config: dict = config.get(DOMAIN) or {}
+
+    def exec_command(call: ServiceCall):
+        host = call.data.get("host", config.get(CONF_HOST, "172.17.0.1"))
+        port = call.data.get("port", config.get(CONF_PORT, 22))
+        username = call.data.get("user", config.get(CONF_USERNAME, "pi"))
+        password = call.data.get("pass", config.get(CONF_PASSWORD, "raspberry"))
+        command = call.data.get("command")
 
         client = SSHClient()
         client.set_missing_host_key_policy(AutoAddPolicy())
 
-        if ssh_private_key_path:
-            try:
-                key = RSAKey.from_private_key_file(ssh_private_key_path)
+        try:
+            if private_key := call.data.get("private_key"):
+                key = RSAKey.from_private_key_file(private_key)
                 client.connect(host, port, username, pkey=key)
-            except Exception as e:
-                _LOGGER.error(f"Failed to connect using SSH key: {e}")
-                return
-        else:
-            # Use password for authentication if SSH key is not provided
-            client.connect(host, port, username, password)
-
-        stdin, stdout, stderr = client.exec_command(command)
-        data = stdout.read()
-        stderr.read()
+            else:
+                # Use password for authentication if SSH key is not provided
+                client.connect(host, port, username, password)
+        except Exception as e:
+            _LOGGER.error(f"Failed to connect: {repr(e)}")
+            return {"error": repr(e)}
+
+        _, stdout, stderr = client.exec_command(command)
+        response = {
+            "command": command,
+            "stdout": stdout.read().decode("utf-8"),
+            "stderr": stderr.read().decode("utf-8"),
+        }
         client.close()
 
-        _LOGGER.info(data)
+        _LOGGER.info(response)
 
-    hass.services.register(DOMAIN, 'exec_command', exec_command)
+    # ServiceResponse from Hass 2023.7
+    # https://github.com/home-assistant/core/blob/2023.7.0/homeassistant/core.py
+    hass.services.async_register(
+        DOMAIN, "exec_command", exec_command, SSH_COMMAND_SCHEMA
+    )
 
     return True
 
 
 async def async_setup_entry(hass, entry):
-    return True
+    return True

custom_components/ssh_command/config_flow.py

@@ -9,5 +9,5 @@ async def async_step_import(self, user_input=None):
 
     async def async_step_user(self, user_input=None):
         if self._async_current_entries():
-            return self.async_abort(reason='single_instance_allowed')
+            return self.async_abort(reason="single_instance_allowed")
         return self.async_create_entry(title="SSH Command", data={})