From 2dafd5ab3f36eaaf13366edb0fa2eaba6cf5524b Mon Sep 17 00:00:00 2001 From: John Dyer Date: Mon, 23 Oct 2023 09:20:53 -0400 Subject: [PATCH 1/3] first draft of larger websocker fix --- homeassistant_cli/cli.py | 23 ++++++++++++++++++++++- homeassistant_cli/config.py | 3 ++- homeassistant_cli/const.py | 2 ++ homeassistant_cli/remote.py | 3 ++- 4 files changed, 28 insertions(+), 3 deletions(-) diff --git a/homeassistant_cli/cli.py b/homeassistant_cli/cli.py index b925eee..46df6f4 100644 --- a/homeassistant_cli/cli.py +++ b/homeassistant_cli/cli.py @@ -23,6 +23,16 @@ Configuration, ensure=True ) +def validate_file_size(ctx, param, value): + if value is not None: + if str(value).endswith('KB'): + size = int(value)[:-2] * 1024 + elif str(value).endswith('MB'): + size = int(value)[:-2] * 1024 * 1024 + else: + # bytes + size = int(value) + return size def run() -> None: """Run entry point. @@ -136,7 +146,7 @@ def _default_token() -> Optional[str]: @click.option( '--timeout', help='Timeout for network operations.', - default=const.DEFAULT_TIMEOUT, + default=str(const.DEFAULT_TIMEOUT), show_default=True, ) @click.option( @@ -205,6 +215,15 @@ def _default_token() -> Optional[str]: default=None, help='Sort table by the jsonpath expression. Example: last_changed', ) + +@click.option( + '--max-message-size', + callback=validate_file_size, + default=const.WS_MAX_MESSAGE_SIZE, + help='Max size of websocket payload. Default: 4MB', + envvar='HASS_WS_MAX_MESSAGE_SIZE', +) + @pass_context def cli( ctx: Configuration, @@ -222,6 +241,7 @@ def cli( no_headers: bool, table_format: str, sort_by: Optional[str], + max_message_size: str, ) -> None: """Command line interface for Home Assistant.""" ctx.verbose = verbose @@ -238,6 +258,7 @@ def cli( ctx.no_headers = no_headers ctx.table_format = table_format ctx.sort_by = sort_by # type: ignore + ctx.max_message_size = max_message_size _LOGGER.debug("Using settings: %s", ctx) diff --git a/homeassistant_cli/config.py b/homeassistant_cli/config.py index 4fd2e43..c980894 100644 --- a/homeassistant_cli/config.py +++ b/homeassistant_cli/config.py @@ -126,7 +126,7 @@ def __init__(self) -> None: self.no_headers = False self.table_format = 'plain' self.sort_by = None - + self.max_message_size = const.WS_MAX_MESSAGE_SIZE # type: str def echo(self, msg: str, *args: Optional[Any]) -> None: """Put content message to stdout.""" self.log(msg, *args) @@ -151,6 +151,7 @@ def __repr__(self) -> str: "access-token": 'yes' if self.token is not None else 'no', "api-password": 'yes' if self.password is not None else 'no', "insecure": self.insecure, + "max_message_size" : self.max_message_size, "output": self.output, "verbose": self.verbose, } diff --git a/homeassistant_cli/const.py b/homeassistant_cli/const.py index bb782d5..a766528 100644 --- a/homeassistant_cli/const.py +++ b/homeassistant_cli/const.py @@ -20,3 +20,5 @@ ('CHANGED', 'last_changed'), ] COLUMNS_SERVICES = [('DOMAIN', 'domain'), ("SERVICE", "domain.services[*]")] + +WS_MAX_MESSAGE_SIZE = '4194304' \ No newline at end of file diff --git a/homeassistant_cli/remote.py b/homeassistant_cli/remote.py index 8f9558b..909c068 100644 --- a/homeassistant_cli/remote.py +++ b/homeassistant_cli/remote.py @@ -105,7 +105,8 @@ def wsapi( async def fetcher() -> Optional[Dict]: async with aiohttp.ClientSession() as session: async with session.ws_connect( - resolve_server(ctx) + "/api/websocket" + resolve_server(ctx) + "/api/websocket", + max_msg_size=ctx.max_message_size ) as wsconn: await wsconn.send_str( From 86ac4f0fd89412fb0e2706f0cf267c86d2691f31 Mon Sep 17 00:00:00 2001 From: John Dyer Date: Fri, 15 Dec 2023 08:26:32 -0500 Subject: [PATCH 2/3] Fix linting --- homeassistant_cli/cli.py | 22 +++++++++++----------- homeassistant_cli/config.py | 7 +++---- homeassistant_cli/const.py | 2 +- homeassistant_cli/remote.py | 3 +-- tests/test_area.py | 2 -- tests/test_device.py | 3 --- tests/test_raw.py | 2 -- tests/test_service.py | 1 - 8 files changed, 16 insertions(+), 26 deletions(-) diff --git a/homeassistant_cli/cli.py b/homeassistant_cli/cli.py index 46df6f4..398daad 100644 --- a/homeassistant_cli/cli.py +++ b/homeassistant_cli/cli.py @@ -23,16 +23,18 @@ Configuration, ensure=True ) + def validate_file_size(ctx, param, value): - if value is not None: - if str(value).endswith('KB'): - size = int(value)[:-2] * 1024 - elif str(value).endswith('MB'): - size = int(value)[:-2] * 1024 * 1024 - else: - # bytes - size = int(value) - return size + if value is not None: + if str(value).endswith('KB'): + size = int(value)[:-2] * 1024 + elif str(value).endswith('MB'): + size = int(value)[:-2] * 1024 * 1024 + else: + # bytes + size = int(value) + return size + def run() -> None: """Run entry point. @@ -215,7 +217,6 @@ def _default_token() -> Optional[str]: default=None, help='Sort table by the jsonpath expression. Example: last_changed', ) - @click.option( '--max-message-size', callback=validate_file_size, @@ -223,7 +224,6 @@ def _default_token() -> Optional[str]: help='Max size of websocket payload. Default: 4MB', envvar='HASS_WS_MAX_MESSAGE_SIZE', ) - @pass_context def cli( ctx: Configuration, diff --git a/homeassistant_cli/config.py b/homeassistant_cli/config.py index c980894..ab6edca 100644 --- a/homeassistant_cli/config.py +++ b/homeassistant_cli/config.py @@ -84,9 +84,7 @@ def resolve_server(ctx: Any) -> str: # noqa: F821 ctx.resolved_server = None if not ctx.resolved_server: - if ctx.server == "auto": - if "HASSIO_TOKEN" in os.environ and "HASS_TOKEN" not in os.environ: ctx.resolved_server = const.DEFAULT_SERVER_MDNS else: @@ -126,7 +124,8 @@ def __init__(self) -> None: self.no_headers = False self.table_format = 'plain' self.sort_by = None - self.max_message_size = const.WS_MAX_MESSAGE_SIZE # type: str + self.max_message_size = const.WS_MAX_MESSAGE_SIZE # type: str + def echo(self, msg: str, *args: Optional[Any]) -> None: """Put content message to stdout.""" self.log(msg, *args) @@ -151,7 +150,7 @@ def __repr__(self) -> str: "access-token": 'yes' if self.token is not None else 'no', "api-password": 'yes' if self.password is not None else 'no', "insecure": self.insecure, - "max_message_size" : self.max_message_size, + "max_message_size": self.max_message_size, "output": self.output, "verbose": self.verbose, } diff --git a/homeassistant_cli/const.py b/homeassistant_cli/const.py index a766528..a91d249 100644 --- a/homeassistant_cli/const.py +++ b/homeassistant_cli/const.py @@ -21,4 +21,4 @@ ] COLUMNS_SERVICES = [('DOMAIN', 'domain'), ("SERVICE", "domain.services[*]")] -WS_MAX_MESSAGE_SIZE = '4194304' \ No newline at end of file +WS_MAX_MESSAGE_SIZE = '4194304' diff --git a/homeassistant_cli/remote.py b/homeassistant_cli/remote.py index 909c068..bcc8fee 100644 --- a/homeassistant_cli/remote.py +++ b/homeassistant_cli/remote.py @@ -106,9 +106,8 @@ async def fetcher() -> Optional[Dict]: async with aiohttp.ClientSession() as session: async with session.ws_connect( resolve_server(ctx) + "/api/websocket", - max_msg_size=ctx.max_message_size + max_msg_size=ctx.max_message_size, ) as wsconn: - await wsconn.send_str( json.dumps({'type': 'auth', 'access_token': ctx.token}) ) diff --git a/tests/test_area.py b/tests/test_area.py index 0750b7b..56a244f 100644 --- a/tests/test_area.py +++ b/tests/test_area.py @@ -13,7 +13,6 @@ def test_area_list(default_areas) -> None: with mock.patch( 'homeassistant_cli.remote.get_areas', return_value=default_areas ): - runner = CliRunner() result = runner.invoke( cli.cli, ["--output=json", "area", "list"], catch_exceptions=False @@ -29,7 +28,6 @@ def test_area_list_filter(default_areas) -> None: with mock.patch( 'homeassistant_cli.remote.get_areas', return_value=default_areas ): - runner = CliRunner() result = runner.invoke( cli.cli, diff --git a/tests/test_device.py b/tests/test_device.py index 3cc63c2..d2ba7d4 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -16,7 +16,6 @@ def test_device_list(default_devices, default_areas) -> None: with mock.patch( 'homeassistant_cli.remote.get_areas', return_value=default_areas ): - runner = CliRunner() result = runner.invoke( cli.cli, @@ -37,7 +36,6 @@ def test_device_list_filter(default_devices, default_areas) -> None: with mock.patch( 'homeassistant_cli.remote.get_areas', return_value=default_areas ): - runner = CliRunner() result = runner.invoke( cli.cli, @@ -64,7 +62,6 @@ def test_device_assign(default_areas, default_devices) -> None: 'homeassistant_cli.remote.assign_area', return_value={'success': True}, ): - runner = CliRunner() result = runner.invoke( cli.cli, diff --git a/tests/test_raw.py b/tests/test_raw.py index f75d70f..4b16882 100644 --- a/tests/test_raw.py +++ b/tests/test_raw.py @@ -82,7 +82,6 @@ def test_raw_ws() -> None: with mocker.patch( 'homeassistant_cli.remote.wsapi', return_value={"result": "worked"} ) as mockmethod: - runner = CliRunner() result = runner.invoke( cli.cli, @@ -103,7 +102,6 @@ def test_raw_ws_data() -> None: with mocker.patch( 'homeassistant_cli.remote.wsapi', return_value={"result": "worked"} ) as mockmethod: - runner = CliRunner() result = runner.invoke( cli.cli, diff --git a/tests/test_service.py b/tests/test_service.py index f180c09..574caef 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -74,7 +74,6 @@ def test_service_completion(default_services) -> None: def test_service_call(default_services) -> None: """Test basic call of a service.""" with requests_mock.Mocker() as mock: - post = mock.post( "http://localhost:8123/api/services/homeassistant/restart", json={"result": "bogus"}, From 5a189964c9b6df314e37a45dfb08749b1053bc7e Mon Sep 17 00:00:00 2001 From: John Dyer Date: Fri, 15 Dec 2023 09:02:42 -0500 Subject: [PATCH 3/3] fix ws size callback function --- homeassistant_cli/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/homeassistant_cli/cli.py b/homeassistant_cli/cli.py index 398daad..029cd56 100644 --- a/homeassistant_cli/cli.py +++ b/homeassistant_cli/cli.py @@ -27,9 +27,9 @@ def validate_file_size(ctx, param, value): if value is not None: if str(value).endswith('KB'): - size = int(value)[:-2] * 1024 + size = int(value[:-2]) * 1024 elif str(value).endswith('MB'): - size = int(value)[:-2] * 1024 * 1024 + size = int(value[:-2]) * 1024 * 1024 else: # bytes size = int(value)