From 8b566757c303f62f2dbaa7845765cadfb3ef421f Mon Sep 17 00:00:00 2001
From: Teemu R <tpr@iki.fi>
Date: Wed, 24 Jan 2024 09:10:55 +0100
Subject: [PATCH] Add new cli command 'command' to execute arbitrary commands
 (#692)

* Add new cli command 'command' to execute arbitrary commands

This deprecates 'raw-command', which requires positional argument for module,
in favor of new 'command' that accepts '--module' option for IOT devices.

* Pull block list to the module level
---
 kasa/cli.py            | 23 +++++++++++++++++++----
 kasa/modules/emeter.py |  2 +-
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/kasa/cli.py b/kasa/cli.py
index 1b20303d..5f726be0 100755
--- a/kasa/cli.py
+++ b/kasa/cli.py
@@ -1,4 +1,5 @@
 """python-kasa cli tool."""
+import ast
 import asyncio
 import json
 import logging
@@ -69,6 +70,9 @@ DEVICE_FAMILY_TYPES = [
     device_family_type.value for device_family_type in DeviceFamilyType
 ]
 
+# Block list of commands which require no update
+SKIP_UPDATE_COMMANDS = ["wifi", "raw-command", "command"]
+
 click.anyio_backend = "asyncio"
 
 pass_dev = click.make_pass_decorator(SmartDevice)
@@ -339,8 +343,9 @@ async def cli(
             credentials=credentials,
         )
 
-    # Skip update for wifi & raw-command, and if factory was used to connect
-    if ctx.invoked_subcommand not in ["wifi", "raw-command"] and not device_family:
+    # Skip update on specific commands, or if device factory,
+    # that performs an update was used for the device.
+    if ctx.invoked_subcommand not in SKIP_UPDATE_COMMANDS and not device_family:
         await dev.update()
 
     ctx.obj = dev
@@ -592,13 +597,23 @@ async def alias(dev, new_alias, index):
 
 @cli.command()
 @pass_dev
+@click.pass_context
 @click.argument("module")
 @click.argument("command")
 @click.argument("parameters", default=None, required=False)
-async def raw_command(dev: SmartDevice, module, command, parameters):
+async def raw_command(ctx, dev: SmartDevice, module, command, parameters):
     """Run a raw command on the device."""
-    import ast
+    logging.warning("Deprecated, use 'kasa command --module %s %s'", module, command)
+    return await ctx.forward(cmd_command)
 
+
+@cli.command(name="command")
+@pass_dev
+@click.option("--module", required=False, help="Module for IOT protocol.")
+@click.argument("command")
+@click.argument("parameters", default=None, required=False)
+async def cmd_command(dev: SmartDevice, module, command, parameters):
+    """Run a raw command on the device."""
     if parameters is not None:
         parameters = ast.literal_eval(parameters)
 
diff --git a/kasa/modules/emeter.py b/kasa/modules/emeter.py
index a205396e..11eed48f 100644
--- a/kasa/modules/emeter.py
+++ b/kasa/modules/emeter.py
@@ -63,7 +63,7 @@ class Emeter(Usage):
         self,
         data: List[Dict[str, Union[int, float]]],
         entry_key: str,
-        kwh: bool=True,
+        kwh: bool = True,
         key: Optional[int] = None,
     ) -> Dict[Union[int, float], Union[int, float]]:
         """Return emeter information keyed with the day/month.