python-kasa · yparitcher · Aug 25, 2023 · Aug 25, 2023 · Aug 25, 2023 · Aug 25, 2023
diff --git a/kasa/cli.py b/kasa/cli.py
@@ -569,6 +569,26 @@
     return res
 
 
+@cli.command()
+@click.option("--list", is_flag=True)
+@click.argument("index", type=int, required=False)
+@pass_dev
+async def set_timezone(dev, list, index):
+    """Set the device timezone.
+
+    pass --list to see valid timezones
+    """
+    time = dev.modules["time"]
+    if list:
+        for timezone in time.get_timezones():
+            echo(
+                f"Index: {timezone['index']:3}\tName: {timezone['zone_str']:65}\tRule: {timezone['tz_str']}"
+            )
+        return
+    if index:
+        return await time.set_timezone(index)
+
+
 @cli.command()
 @click.option("--index", type=int, required=False)
 @click.option("--name", type=str, required=False)
@@ -630,18 +650,136 @@
 
 @schedule.command(name="list")
 @pass_dev
+@click.option("--json", is_flag=True)
 @click.argument("type", default="schedule")
-def _schedule_list(dev, type):
+def _schedule_list(dev, json, type):
     """Return the list of schedule actions for the given type."""
     sched = dev.modules[type]
-    for rule in sched.rules:
-        print(rule)
-    else:
+    if sched.rules == []:
         echo(f"No rules of type {type}")
+    for rule in sched.rules:
+        if json:
+            print(rule.json())
+        else:
+            print(rule)
 
     return sched.rules
 
 
+@schedule.command(name="enable")
+@pass_dev
+@click.argument("enable", type=click.BOOL)
+async def _schedule_enable(dev, enable):
+    """Enable or disable schedule."""
+    schedule = dev.modules["schedule"]
+    return await schedule.set_enabled(1 if state else 0)
+
+
+@schedule.command(name="add")
+@pass_dev
+@click.option("--name", type=str, required=True)
+@click.option("--enable", type=click.BOOL, default=True, show_default=True)
+@click.option("--repeat", type=click.BOOL, default=True, show_default=True)
+@click.option("--days", type=str, required=True)
+@click.option("--start-action", type=click.IntRange(-1, 2), default=None, required=True)
+@click.option("--start-sun", type=click.IntRange(-1, 2), default=None, required=True)
+@click.option(
+    "--start-minutes", type=click.IntRange(0, 1440), default=None, required=True
+)
+@click.option("--end-action", type=click.IntRange(-1, 2), default=-1)
+@click.option("--end-sun", type=click.IntRange(-1, 2), default=-1)
+@click.option("--end-minutes", type=click.IntRange(0, 1440), default=None)
+async def add_rule(
+    dev,
+    name,
+    enable,
+    repeat,
+    days,
+    start_action,
+    start_sun,
+    start_minutes,
+    end_action,
+    end_sun,
+    end_minutes,
+):
+    """Add rule to device."""
+    schedule = dev.modules["schedule"]
+    rule_to_add = schedule.Rule(
+        name=name,
+        enable=enable,
+        repeat=repeat,
+        days=list(map(int, days.split(","))),
+        start_action=start_action,
+        start_sun=start_sun,
+        start_minutes=start_minutes,
+        end_action=end_action,
+        end_sun=end_sun,
+        end_minutes=end_minutes,
+    )
+    if rule_to_add:
+        echo("Adding rule")
+        return await schedule.add_rule(rule_to_add)
+    else:
+        echo("Invalid rule")
+
+
+@schedule.command(name="edit")
+@pass_dev
+@click.option("--id", type=str, required=True)
+@click.option("--name", type=str)
+@click.option("--enable", type=click.BOOL)
+@click.option("--repeat", type=click.BOOL)
+@click.option("--days", type=str)
+@click.option("--start-action", type=click.IntRange(-1, 2))
+@click.option("--start-sun", type=click.IntRange(-1, 2))
+@click.option("--start-minutes", type=click.IntRange(0, 1440))
+@click.option("--end-action", type=click.IntRange(-1, 2))
+@click.option("--end-sun", type=click.IntRange(-1, 2))
+@click.option("--end-minutes", type=click.IntRange(0, 1440))
+async def edit_rule(
+    dev,
+    id,
+    name,
+    enable,
+    repeat,
+    days,
+    start_action,
+    start_sun,
+    start_minutes,
+    end_action,
+    end_sun,
+    end_minutes,
+):
+    """Edit rule from device."""
+    schedule = dev.modules["schedule"]
+    rule_to_edit = next(filter(lambda rule: (rule.id == id), schedule.rules), None)
+    if rule_to_edit:
+        echo(f"Editing rule id {id}")
+        if name is not None:
+            rule_to_edit.name = name
+        if enable is not None:
+            rule_to_edit.enable = 1 if enable else 0
+        if repeat is not None:
+            rule_to_edit.repeat = 1 if repeat else 0
+        if days is not None:
+            rule_to_edit.wday = list(map(int, days.split(",")))
+        if start_action is not None:
+            rule_to_edit.sact = start_action
+        if start_sun is not None:
+            rule_to_edit.stime_opt = start_sun
+        if start_minutes is not None:
+            rule_to_edit.smin = start_minutes
+        if end_action is not None:
+            rule_to_edit.eact = end_action
+        if end_sun is not None:
+            rule_to_edit.etime_opt = end_sun
+        if end_minutes is not None:
+            rule_to_edit.emin = end_minutes
+        return await schedule.edit_rule(rule_to_edit)
+    else:
+        echo(f"No rule with id {id} was found")
+
+
 @schedule.command(name="delete")
 @pass_dev
 @click.option("--id", type=str, required=True)
@@ -656,6 +794,17 @@
         echo(f"No rule with id {id} was found")
 
 
+@schedule.command()
+@pass_dev
+@click.option("--prompt", type=click.BOOL, prompt=True, help="Are you sure?")
+async def delete_all(dev, prompt):
+    """Delete all rules from device."""
+    schedule = dev.modules["schedule"]
+    if prompt:
+        echo("Deleting all rules")
+        return await schedule.delete_all_rules()
+
+
 @cli.group(invoke_without_command=True)
 @click.pass_context
 async def presets(ctx):

diff --git a/kasa/modules/__init__.py b/kasa/modules/__init__.py
@@ -6,8 +6,14 @@
 from .emeter import Emeter
 from .module import Module
 from .motion import Motion
-from .rulemodule import Rule, RuleModule
-from .schedule import Schedule
+from .rulemodule import (
+    AntitheftRule,
+    BulbScheduleRule,
+    CountdownRule,
+    RuleModule,
+    ScheduleRule,
+)
+from .schedule import BulbSchedule, Schedule
 from .time import Time
 from .usage import Usage
 
@@ -19,9 +25,13 @@
     "Emeter",
     "Module",
     "Motion",
-    "Rule",
+    "AntitheftRule",
+    "CountdownRule",
+    "ScheduleRule",
+    "BulbScheduleRule",
     "RuleModule",
     "Schedule",
+    "BulbSchedule",
     "Time",
     "Usage",
 ]
diff --git a/kasa/modules/antitheft.py b/kasa/modules/antitheft.py
@@ -1,9 +1,11 @@
 """Implementation of the antitheft module."""
-from .rulemodule import RuleModule
+from .rulemodule import AntitheftRule, RuleModule
 
 
 class Antitheft(RuleModule):
     """Implementation of the antitheft module.
 
     This shares the functionality among other rule-based modules.
     """
+
+    Rule = AntitheftRule
diff --git a/kasa/modules/countdown.py b/kasa/modules/countdown.py
@@ -1,6 +1,8 @@
 """Implementation for the countdown timer."""
-from .rulemodule import RuleModule
+from .rulemodule import CountdownRule, RuleModule
 
 
 class Countdown(RuleModule):
     """Implementation of countdown module."""
+
+    Rule = CountdownRule
diff --git a/kasa/modules/rulemodule.py b/kasa/modules/rulemodule.py
@@ -1,13 +1,21 @@
 """Base implementation for all rule-based modules."""
+import json
 import logging
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from .module import Module, merge
 
 
+class EnabledOption(Enum):
+    """Integer enabled option."""
+
+    TurnOff = 0
+    Enabled = 1
+
+
 class Action(Enum):
     """Action to perform."""
 
@@ -26,55 +34,114 @@
     AtSunset = 2
 
 
-class Rule(BaseModel):
+class BulbModeOption(Enum):
+    """Bulb mode."""
+
+    Customize = "customize_preset"
+    Last = "last_status"
+
+
+class BaseRule(BaseModel):
     """Representation of a rule."""
 
-    id: str
+    # not used when adding a rule
+    id: Optional[str]
     name: str
-    enable: bool
-    wday: List[int]
-    repeat: bool
+    enable: EnabledOption
+
+    class Config:
+        """Rule Config."""
+
+        validate_assignment = True
+        allow_population_by_field_name = True
+
+
+class CountdownRule(BaseRule):
+    """Representation of a countdown rule."""
+
+    delay: int
+    act: EnabledOption = Field(alias="action")
+
+
+class ScheduleRule(BaseRule):
+    """Representation of a schedule rule."""
+
+    wday: List[int] = Field(alias="days")
+    repeat: EnabledOption
 
     # start action
-    sact: Optional[Action]
-    stime_opt: TimeOption
-    smin: int
+    sact: Optional[Action] = Field(alias="start_action")
-    sact: Optional[Action] = Field(alias="start_action")
+    start_action: Optional[Action] = Field(alias="sact")
-    sact: Optional[Action] = Field(alias="start_action")
+    start_action: Optional[Action] = Field(alias="sact")
+    stime_opt: TimeOption = Field(alias="start_sun")
+    smin: int = Field(alias="start_minutes", ge=0, le=1440)
+
+    # end action
+    eact: Optional[Action] = Field(alias="end_action")
+    # Required to submit, but the device will not return it if set to -1
+    etime_opt: TimeOption = Field(default=TimeOption.Disabled, alias="end_sun")
+    emin: Optional[int] = Field(alias="end_minutes", ge=0, le=1440)
+
+
+class BulbRule(BaseRule):
+    """Representation of a bulb schedule rule."""
+
+    saturation: int = Field(ge=0, le=100)
+    hue: int = Field(ge=0, le=360)
+    brightness: int = Field(ge=0, le=100)
+    color_temp: int = Field(ge=2500, le=9000)
+    mode: BulbModeOption
+    on_off: EnabledOption
+
 
-    eact: Optional[Action]
-    etime_opt: TimeOption
-    emin: int
+class BulbScheduleRule(BaseRule):
+    """Representation of a bulb schedule rule."""
 
-    # Only on bulbs
-    s_light: Optional[Dict]
+    s_light: BulbRule = Field(alias="lights")
+
+
+class AntitheftRule(BaseRule):
+    """Representation of a antitheft rule."""
+
+    frequency: int = Field(ge=1, le=10)
 
 
 _LOGGER = logging.getLogger(__name__)
 
 
 class RuleModule(Module):
-    """Base class for rule-based modules, such as countdown and antitheft."""
+    """Base class for rule-based modules, such as antitheft, countdown and schedule."""
+
+    Rule = BaseRule
 
     def query(self):
         """Prepare the query for rules."""
         q = self.query_for_command("get_rules")
         return merge(q, self.query_for_command("get_next_action"))
 
     @property
-    def rules(self) -> List[Rule]:
+    def rules(self) -> List[BaseRule]:
         """Return the list of rules for the service."""
         try:
             return [
-                Rule.parse_obj(rule) for rule in self.data["get_rules"]["rule_list"]
+                self.Rule.parse_obj(rule)
+                for rule in self.data["get_rules"]["rule_list"]
             ]
         except Exception as ex:
             _LOGGER.error("Unable to read rule list: %s (data: %s)", ex, self.data)
             return []
 
-    async def set_enabled(self, state: bool):
+    async def set_enabled(self, state: int):
         """Enable or disable the service."""
-        return await self.call("set_overall_enable", state)
+        return await self.call("set_overall_enable", {"enable": state})
+
+    async def add_rule(self, rule: BaseRule):
+        """Add a new rule."""
+        return await self.call("add_rule", json.loads(rule.json(exclude_none=True)))
+
+    async def edit_rule(self, rule: BaseRule):
+        """Edit the given rule."""
+        return await self.call("edit_rule", json.loads(rule.json(exclude_none=True)))
 
-    async def delete_rule(self, rule: Rule):
+    async def delete_rule(self, rule: BaseRule):
         """Delete the given rule."""
         return await self.call("delete_rule", {"id": rule.id})