relay_mqtt_bot, pump_mqtt_bot

7 files changed, 212 insertions, 56 deletions

diff --git a/src/home/api/types/types.py b/src/home/api/types/types.py
index 3f77b19..981e798 100644
--- a/src/home/api/types/types.py
+++ b/src/home/api/types/types.py
@@ -9,6 +9,7 @@ class BotType(Enum):
     SOUND = auto()
     POLARIS_KETTLE = auto()
     PUMP_MQTT = auto()
+    RELAY_MQTT = auto()
 
 
 class TemperatureSensorLocation(Enum):
diff --git a/src/home/mqtt/__init__.py b/src/home/mqtt/__init__.py
index a6f5f5e..c9a6c6e 100644
--- a/src/home/mqtt/__init__.py
+++ b/src/home/mqtt/__init__.py
@@ -1,3 +1,3 @@
 from .mqtt import MQTTBase
 from .util import poll_tick
-from .relay import MQTTRelay
\ No newline at end of file
+from .relay import MQTTRelay, MQTTRelayState, MQTTRelayDevice
\ No newline at end of file
diff --git a/src/home/mqtt/relay.py b/src/home/mqtt/relay.py
index 0f97b5b..b481bf8 100644
--- a/src/home/mqtt/relay.py
+++ b/src/home/mqtt/relay.py
@@ -1,8 +1,9 @@
 import paho.mqtt.client as mqtt
 import re
+import datetime
 
 from .mqtt import MQTTBase
-from typing import Optional, Union
+from typing import Optional, Union, List
 from .payload.relay import (
     InitialStatPayload,
     StatPayload,
@@ -11,19 +12,27 @@ from .payload.relay import (
 )
 
 
+class MQTTRelayDevice:
+    home_id: str
+    secret: str
+
+    def __init__(self, home_id: str, secret: str):
+        self.home_id = home_id
+        self.secret = secret
+
+
 class MQTTRelay(MQTTBase):
-    _home_id: Union[str, int]
-    _secret: str
+    _devices: list[MQTTRelayDevice]
     _message_callback: Optional[callable]
     _ota_publish_callback: Optional[callable]
 
     def __init__(self,
-                 home_id: Union[str, int],
-                 secret: str,
+                 devices: Union[MQTTRelayDevice, list[MQTTRelayDevice]],
                  subscribe_to_updates=True):
         super().__init__(clean_session=True)
-        self._home_id = home_id
-        self._secret = secret
+        if not isinstance(devices, list):
+            devices = [devices]
+        self._devices = devices
         self._message_callback = None
         self._ota_publish_callback = None
         self._subscribe_to_updates = subscribe_to_updates
@@ -33,10 +42,10 @@ class MQTTRelay(MQTTBase):
         super().on_connect(client, userdata, flags, rc)
 
         if self._subscribe_to_updates:
-            topic = f'hk/{self._home_id}/relay/#'
-            self._logger.info(f"subscribing to {topic}")
-
-            client.subscribe(topic, qos=1)
+            for device in self._devices:
+                topic = f'hk/{device.home_id}/relay/#'
+                self._logger.info(f"subscribing to {topic}")
+                client.subscribe(topic, qos=1)
 
     def on_publish(self, client: mqtt.Client, userdata, mid):
         if self._ota_mid is not None and mid == self._ota_mid and self._ota_publish_callback:
@@ -55,7 +64,7 @@ class MQTTRelay(MQTTBase):
             name = match.group(1)
             subtopic = match.group(2)
 
-            if name != self._home_id:
+            if name not in self._devices:
                 return
 
             message = None
@@ -67,27 +76,56 @@ class MQTTRelay(MQTTBase):
                 message = PowerPayload.unpack(msg.payload)
 
             if message and self._message_callback:
-                self._message_callback(message)
+                self._message_callback(name, message)
 
         except Exception as e:
             self._logger.exception(str(e))
 
-    def set_power(self, enable: bool):
-        payload = PowerPayload(secret=self._secret,
+    def set_power(self, home_id, enable: bool):
+        device = next(d for d in self._devices if d.home_id == home_id)
+
+        payload = PowerPayload(secret=device.secret,
                                state=enable)
-        self._client.publish(f'hk/{self._home_id}/relay/power',
+        self._client.publish(f'hk/{device.home_id}/relay/power',
                              payload=payload.pack(),
                              qos=1)
         self._client.loop_write()
 
     def push_ota(self,
+                 home_id,
                  filename: str,
                  publish_callback: callable,
                  qos: int):
+        device = next(d for d in self._devices if d.home_id == home_id)
+
         self._ota_publish_callback = publish_callback
-        payload = OTAPayload(secret=self._secret, filename=filename)
-        publish_result = self._client.publish(f'hk/{self._home_id}/relay/admin/ota',
+        payload = OTAPayload(secret=device.secret, filename=filename)
+        publish_result = self._client.publish(f'hk/{device.home_id}/relay/admin/ota',
                                               payload=payload.pack(),
                                               qos=qos)
         self._ota_mid = publish_result.mid
         self._client.loop_write()
+
+
+class MQTTRelayState:
+    enabled: bool
+    update_time: datetime.datetime
+    rssi: int
+    fw_version: int
+    ever_updated: bool
+
+    def __init__(self):
+        self.ever_updated = False
+        self.enabled = False
+        self.rssi = 0
+
+    def update(self,
+               enabled: bool,
+               rssi: int,
+               fw_version=None):
+        self.ever_updated = True
+        self.enabled = enabled
+        self.rssi = rssi
+        self.update_time = datetime.datetime.now()
+        if fw_version:
+            self.fw_version = fw_version
diff --git a/src/home/telegram/bot.py b/src/home/telegram/bot.py
index bd09f42..9e60b70 100644
--- a/src/home/telegram/bot.py
+++ b/src/home/telegram/bot.py
@@ -123,15 +123,25 @@ def handler(**kwargs):
             return _handler_of_handler(f=f, *args, **inner_kwargs)
 
         messages = []
+        texts = []
+
         if 'messages' in kwargs:
             messages += kwargs['messages']
         if 'message' in kwargs:
             messages.append(kwargs['message'])
+
+        if 'text' in kwargs:
+            texts.append(kwargs['text'])
+        if 'texts' in kwargs:
+            texts.append(kwargs['texts'])
+
         if messages:
-            _updater.dispatcher.add_handler(
-                MessageHandler(text_filter(*list(itertools.chain.from_iterable([lang.all(m) for m in messages]))), _handler),
-                group=0
-            )
+            texts = list(itertools.chain.from_iterable([lang.all(m) for m in messages]))
+
+        _updater.dispatcher.add_handler(
+            MessageHandler(text_filter(*texts), _handler),
+            group=0
+        )
 
         if 'command' in kwargs:
             _updater.dispatcher.add_handler(CommandHandler(kwargs['command'], _handler), group=0)
diff --git a/src/pump_mqtt_bot.py b/src/pump_mqtt_bot.py
index d87234b..f16ed36 100755
--- a/src/pump_mqtt_bot.py
+++ b/src/pump_mqtt_bot.py
@@ -9,7 +9,7 @@ from home.config import config
 from home.telegram import bot
 from home.telegram._botutil import user_any_name
 from home.api.types import BotType
-from home.mqtt import MQTTRelay
+from home.mqtt import MQTTRelay, MQTTRelayState, MQTTRelayDevice
 from home.mqtt.payload import MQTTPayload
 from home.mqtt.payload.relay import InitialStatPayload, StatPayload
 
@@ -70,30 +70,8 @@ bot.lang.en(
 )
 
 
-class RelayState:
-    enabled: bool
-    update_time: datetime.datetime
-    rssi: int
-    fw_version: int
-    ever_updated: bool
-
-    def __init__(self):
-        self.ever_updated = False
-
-    def update(self,
-               enabled: bool,
-               rssi: int,
-               fw_version=None):
-        self.ever_updated = True
-        self.enabled = enabled
-        self.rssi = rssi
-        self.update_time = datetime.datetime.now()
-        if fw_version:
-            self.fw_version = fw_version
-
-
 mqtt_relay: Optional[MQTTRelay] = None
-relay_state = RelayState()
+relay_state = MQTTRelayState()
 
 
 class UserAction(Enum):
@@ -101,7 +79,7 @@ class UserAction(Enum):
     OFF = 'off'
 
 
-def on_mqtt_message(message: MQTTPayload):
+def on_mqtt_message(home_id, message: MQTTPayload):
     if isinstance(message, InitialStatPayload) or isinstance(message, StatPayload):
         kwargs = dict(rssi=message.rssi, enabled=message.flags.state)
         if isinstance(message, InitialStatPayload):
@@ -121,14 +99,14 @@ def notify(user: User, action: UserAction) -> None:
 
 @bot.handler(message='enable')
 def enable_handler(ctx: bot.Context) -> None:
-    mqtt_relay.set_power(True)
+    mqtt_relay.set_power(config['mqtt']['home_id'], True)
     ctx.reply(ctx.lang('done'))
     notify(ctx.user, UserAction.ON)
 
 
 @bot.handler(message='disable')
 def disable_handler(ctx: bot.Context) -> None:
-    mqtt_relay.set_power(False)
+    mqtt_relay.set_power(config['mqtt']['home_id'], False)
     ctx.reply(ctx.lang('done'))
     notify(ctx.user, UserAction.OFF)
 
@@ -179,8 +157,8 @@ def markup(ctx: Optional[bot.Context]) -> Optional[ReplyKeyboardMarkup]:
 
 
 if __name__ == '__main__':
-    mqtt_relay = MQTTRelay(home_id=config['mqtt']['home_id'],
-                           secret=config['mqtt']['relay']['secret'])
+    mqtt_relay = MQTTRelay(devices=MQTTRelayDevice(home_id=config['mqtt']['home_id'],
+                                                   secret=config['mqtt']['home_secret']))
     mqtt_relay.set_message_callback(on_mqtt_message)
     mqtt_relay.configure_tls()
     mqtt_relay.connect_and_loop(loop_forever=False)
diff --git a/src/relay_mqtt_bot.py b/src/relay_mqtt_bot.py
new file mode 100755
index 0000000..ff24924
--- /dev/null
+++ b/src/relay_mqtt_bot.py
@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+from enum import Enum
+from typing import Optional
+from telegram import ReplyKeyboardMarkup
+from functools import partial
+
+from home.config import config
+from home.telegram import bot
+from home.api.types import BotType
+from home.mqtt import MQTTRelay, MQTTRelayState, MQTTRelayDevice
+from home.mqtt.payload import MQTTPayload
+from home.mqtt.payload.relay import InitialStatPayload, StatPayload
+
+
+config.load('relay_mqtt_bot')
+
+bot.initialize()
+bot.lang.ru(
+    start_message="Выберите команду на клавиатуре",
+    unknown_command="Неизвестная команда",
+
+    enable="Включить",
+    enabled="Включен ✅",
+    disable="Выключить",
+    disabled="Выключен ❌",
+
+    status="Статус",
+    status_updated=' (обновлено %s)',
+
+    done="Готово 👌",
+)
+bot.lang.en(
+    start_message="Select command on the keyboard",
+    unknown_command="Unknown command",
+
+    enable="Turn ON",
+    enabled="Turned ON ✅",
+    disable="Turn OFF",
+    disabled="Turned OFF ❌",
+
+    status="Status",
+    status_updated=' (updated %s)',
+
+    done="Done 👌",
+)
+
+
+type_emojis = {
+    'lamp': '💡'
+}
+status_emoji = {
+    'on': '✅',
+    'off': '❌'
+}
+mqtt_relay: Optional[MQTTRelay] = None
+relay_states: dict[str, MQTTRelayState] = {}
+
+
+class UserAction(Enum):
+    ON = 'on'
+    OFF = 'off'
+
+
+def on_mqtt_message(home_id, message: MQTTPayload):
+    if isinstance(message, InitialStatPayload) or isinstance(message, StatPayload):
+        kwargs = dict(rssi=message.rssi, enabled=message.flags.state)
+        if isinstance(message, InitialStatPayload):
+            kwargs['fw_version'] = message.fw_version
+        if home_id not in relay_states[home_id]:
+            relay_states[home_id] = MQTTRelayState()
+        relay_states[home_id].update(**kwargs)
+
+
+def enable_handler(home_id: str, ctx: bot.Context) -> None:
+    mqtt_relay.set_power(home_id, True)
+    ctx.reply(ctx.lang('done'))
+
+
+def disable_handler(home_id: str, ctx: bot.Context) -> None:
+    mqtt_relay.set_power(home_id, False)
+    ctx.reply(ctx.lang('done'))
+
+
+def start(ctx: bot.Context) -> None:
+    ctx.reply(ctx.lang('start_message'))
+
+
+@bot.exceptionhandler
+def exception_handler(e: Exception, ctx: bot.Context) -> bool:
+    return False
+
+
+@bot.defaultreplymarkup
+def markup(ctx: Optional[bot.Context]) -> Optional[ReplyKeyboardMarkup]:
+    buttons = []
+    for device_id, data in config['relays'].items():
+        labels = data['labels']
+        type_emoji = type_emojis[data['type']]
+        row = [f'{type_emoji}{status_emoji[i.value]} {labels[ctx.user_lang]}'
+               for i in UserAction]
+        buttons.append(row)
+    return ReplyKeyboardMarkup(buttons, one_time_keyboard=False)
+
+
+if __name__ == '__main__':
+    devices = []
+    for device_id, data in config['relays'].items():
+        devices.append(MQTTRelayDevice(home_id=device_id,
+                                       secret=data['secret']))
+        labels = data['labels']
+        bot.lang.ru(**{device_id: labels['ru']})
+        bot.lang.en(**{device_id: labels['en']})
+
+        type_emoji = type_emojis[data['type']]
+
+        for action in UserAction:
+            messages = []
+            for _lang, _label in labels.items():
+                messages.append(f'{type_emoji}{status_emoji[action.value]} {labels[_lang]}')
+            bot.handler(texts=messages)(partial(enable_handler if action == UserAction.ON else disable_handler, device_id))
+
+    mqtt_relay = MQTTRelay(devices=devices)
+    mqtt_relay.set_message_callback(on_mqtt_message)
+    mqtt_relay.configure_tls()
+    mqtt_relay.connect_and_loop(loop_forever=False)
+
+    # bot.enable_logging(BotType.RELAY_MQTT)
+    bot.run(start_handler=start)
+
+    mqtt_relay.disconnect()
diff --git a/tools/mcuota.py b/tools/mcuota.py
index f7f2968..8d47c8c 100755
--- a/tools/mcuota.py
+++ b/tools/mcuota.py
@@ -10,7 +10,7 @@ sys.path.extend([
 from time import sleep
 from argparse import ArgumentParser
 from src.home.config import config
-from src.home.mqtt import MQTTRelay
+from src.home.mqtt import MQTTRelay, MQTTRelayDevice
 
 
 def guess_filename(product: str, build_target: str):
@@ -34,11 +34,10 @@ def relayctl_publish_ota(filename: str,
         global stop
         stop = True
 
-    mqtt_relay = MQTTRelay(home_id=home_id,
-                           secret=home_secret)
+    mqtt_relay = MQTTRelay(devices=MQTTRelayDevice(home_id=home_id, secret=home_secret))
     mqtt_relay.configure_tls()
     mqtt_relay.connect_and_loop(loop_forever=False)
-    mqtt_relay.push_ota(filename, published, qos)
+    mqtt_relay.push_ota(home_id, filename, published, qos)
     while not stop:
         sleep(0.1)
     mqtt_relay.disconnect()