multiple fixes

3 files changed, 58 insertions, 26 deletions

diff --git a/src/home/mqtt/_node.py b/src/home/mqtt/_node.py
index 9e39911..f34da0c 100644
--- a/src/home/mqtt/_node.py
+++ b/src/home/mqtt/_node.py
@@ -1,7 +1,7 @@
 import paho.mqtt.client as mqtt
 
 from .mqtt import MqttBase
-from typing import List
+from typing import List, Optional
 from ._module import MqttModule
 from ._payload import MqttPayload
 
@@ -92,7 +92,10 @@ class MqttNode(MqttBase):
         self._module_subscriptions[topic] = module
         self._client.subscribe(f'hk/{self._node_id}/{topic}', qos)
 
-    def publish(self, topic: str, payload: bytes, qos: int = 1):
+    def publish(self,
+                topic: str,
+                payload: bytes,
+                qos: int = 1):
         self._client.publish(f'hk/{self._node_id}/{topic}', payload, qos)
         self._client.loop_write()
 
diff --git a/src/home/mqtt/module/ota.py b/src/home/mqtt/module/ota.py
index 86d6839..5a1a309 100644
--- a/src/home/mqtt/module/ota.py
+++ b/src/home/mqtt/module/ota.py
@@ -41,27 +41,41 @@ class OtaPayload(MqttPayload):
 
 
 class MqttOtaModule(MqttModule):
+    _ota_request: Optional[tuple[str, str, int]]
+    _mqtt_ref: Optional[MqttNode]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._ota_request = None
+        self._mqtt_ref = None
+
     def init(self, mqtt: MqttNode):
         mqtt.subscribe_module("otares", self)
 
+        self._mqtt_ref = mqtt
+
+        if self._ota_request is not None:
+            secret, filename, qos = self._ota_request
+            self._ota_request = None
+            self.do_push_ota(secret, filename, qos)
+
     def handle_payload(self, mqtt: MqttNode, topic: str, payload: bytes) -> Optional[MqttPayload]:
         if topic == 'otares':
             message = OtaResultPayload.unpack(payload)
             self._logger.debug(message)
             return message
 
-    # def push_ota(self,
-    #              node_id,
-    #              filename: str,
-    #              publish_callback: callable,
-    #              qos: int):
-    #     device = next(d for d in self._devices if d.id == device_id)
-    #     assert device.secret is not None, 'device secret not specified'
-    #
-    #     self._ota_publish_callback = publish_callback
-    #     payload = OtaPayload(secret=device.secret, filename=filename)
-    #     publish_result = self._client.publish(f'hk/{device.id}/{self.TOPIC_LEAF}/admin/ota',
-    #                                           payload=payload.pack(),
-    #                                           qos=qos)
-    #     self._ota_mid = publish_result.mid
-    #     self._client.loop_write()
\ No newline at end of file
+    def do_push_ota(self, secret: str, filename: str, qos: int):
+        payload = OtaPayload(secret=secret, filename=filename)
+        self._mqtt_ref.publish('ota',
+                               payload=payload.pack(),
+                               qos=qos)
+
+    def push_ota(self,
+                 secret: str,
+                 filename: str,
+                 qos: int):
+        if not self._initialized:
+            self._ota_request = (secret, filename, qos)
+        else:
+            self.do_push_ota(secret, filename, qos)
\ No newline at end of file
diff --git a/src/mqtt_node_util.py b/src/mqtt_node_util.py
index 674b60c..0352c8f 100755
--- a/src/mqtt_node_util.py
+++ b/src/mqtt_node_util.py
@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+import os.path
+from time import sleep
 from typing import Optional
 from argparse import ArgumentParser, ArgumentError
 
@@ -25,34 +27,47 @@ if __name__ == '__main__':
                         help='mqtt modules to include')
     parser.add_argument('--switch-relay', choices=[0, 1], type=int,
                         help='send relay state')
-    parser.add_argument('--switch-relay-secret', type=str,
-                        help='secret password to switch relay')
+    parser.add_argument('--push-ota', type=str, metavar='OTA_FILENAME',
+                        help='push ota, argument receives filename')
+    parser.add_argument('--node-secret', type=str,
+                        help='node admin password')
 
     config.load('mqtt_util', parser=parser)
     arg = parser.parse_args()
 
-    if (arg.switch_relay is not None or arg.switch_relay_secret is not None) and 'relay' not in arg.modules:
+    if (arg.switch_relay is not None or arg.node_secret is not None) and 'relay' not in arg.modules:
         raise ArgumentError(None, '--relay is only allowed when \'relay\' module included in --modules')
 
-    if (arg.switch_relay is not None and arg.switch_relay_secret is None) or (arg.switch_relay is None and arg.switch_relay_secret is not None):
-        raise ArgumentError(None, 'both --switch-relay and --switch-relay-secret are required')
-
     mqtt = MqttNode(node_id=arg.node_id)
 
     # must-have modules
-    add_module('ota')
+    ota_module = add_module('ota')
     add_module('diagnostics')
 
     if arg.modules:
         for m in arg.modules:
             module_instance = add_module(m)
             if m == 'relay' and arg.switch_relay is not None:
+                if not arg.node_secret:
+                    raise ArgumentError(None, '--switch-relay requires --node-secret')
                 module_instance.switchpower(mqtt,
                                             arg.switch_relay == 1,
-                                            arg.switch_relay_secret)
+                                            arg.node_secret)
 
     mqtt.configure_tls()
     try:
-        mqtt.connect_and_loop()
+        mqtt.connect_and_loop(loop_forever=False)
+
+        if arg.push_ota:
+            if not os.path.exists(arg.push_ota):
+                raise OSError(f'--push-ota: file \"{arg.push_ota}\" does not exists')
+        if not arg.node_secret:
+            raise ArgumentError(None, 'pushing OTA requires --node-secret')
+
+        ota_module.push_ota(arg.node_secret, arg.push_ota, 1)
+
+        while True:
+            sleep(0.1)
+
     except KeyboardInterrupt:
         mqtt.disconnect()