feat: initial services definition

Author: KC3ZVD

src/kc3zvd/iot_state/listeners/__init__.py

@@ -0,0 +1,104 @@
+from zeroconf import ServiceStateChange, Zeroconf
+from typing import Optional, cast
+from iot_state.workers import wled
+import logging
+import asyncio
+import sys
+import os
+
+from zeroconf.asyncio import (
+    AsyncServiceBrowser,
+    AsyncServiceInfo,
+    AsyncZeroconf,
+)
+
+_PENDING_TASKS: set[asyncio.Task] = set()
+_WLED_SUPPORT = "0.15.0"
+CELERY_BROKER_URL = os.getenv("CELERY_BROKER_URL")
+CELERY_BACKEND_URL = os.getenv("CELERY_BACKEND_URL")
+CELERY_RESULTS_URL = os.getenv("CELERY_RESULTS_URL")
+
+logger = logging.getLogger(__name__)
+
+
+def async_on_service_state_change(
+    zeroconf: Zeroconf, service_type: str, name: str, state_change: ServiceStateChange
+) -> None:
+    if not name.startswith('wled'):
+      logger.debug(f"Service {name} does not appear to be a WLED service, skipping...")
+      return
+    if state_change is not ServiceStateChange.Added:
+      logger.debug(f"State Change: {state_change}")
+      return
+    task = asyncio.ensure_future(async_publish_service_info(zeroconf, service_type, name))
+    _PENDING_TASKS.add(task)
+    task.add_done_callback(_PENDING_TASKS.discard)
+
+
+async def async_publish_service_info(zeroconf: Zeroconf, service_type: str, name: str) -> None:
+    info = AsyncServiceInfo(service_type, name)
+    await info.async_request(zeroconf, 3000)
+    logger.debug("Info from zeroconf.get_service_info: %r" % (info))
+    if info:
+        addresses = ["%s:%d" % (addr, cast(int, info.port)) for addr in info.parsed_scoped_addresses()]
+
+      
+        details = {
+          "source": 'mdns',
+          "platform": 'wled',
+          "attributes": {
+            "addresses": addresses
+          }
+        }
+
+        wled.discover.delay(details=details)
+
+        logger.info(f"Adding service {name} to MQTT")
+    else:
+        logger.warning("No service info available for %s, skipping..." % (name))
+
+class AsyncRunner:
+    def __init__(self) -> None:
+        self.aiobrowser: Optional[AsyncServiceBrowser] = None
+        self.aiozc: Optional[AsyncZeroconf] = None
+
+    async def async_run(self) -> None:
+        self.aiozc = AsyncZeroconf()
+
+        services = ["_http._tcp.local."]
+        logger.debug("Watching for %s service(s)" % services)
+        logger.info("Monitoring mDNS...")
+        
+        self.aiobrowser = AsyncServiceBrowser(
+            self.aiozc.zeroconf, services, handlers=[async_on_service_state_change]
+        )
+        while True:
+            await asyncio.sleep(1)
+
+    async def async_close(self) -> None:
+        assert self.aiozc is not None
+        assert self.aiobrowser is not None
+        await self.aiobrowser.async_cancel()
+        await self.aiozc.async_close()
+
+def listen() -> None:
+
+  # Set up logging
+  logger.setLevel(level=logging.DEBUG)
+  handler = logging.StreamHandler(sys.stdout)
+  handler.setLevel(level=logging.DEBUG)
+  formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+  handler.setFormatter(formatter)
+  logger.addHandler(handler)
+
+  # set up event loop
+  loop = asyncio.get_event_loop()
+  runner = AsyncRunner()
+
+  try:
+    loop.run_until_complete(runner.async_run())
+  except KeyboardInterrupt:
+    loop.run_until_complete(runner.async_close())
+
+if __name__ == "__main__":
+   listen()

src/kc3zvd/iot_state/listeners/wled.py

@@ -0,0 +1 @@
+

src/kc3zvd/iot_state/wled/__init__.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+import paho.mqtt.publish as publish
+import json
+import os
+
+APP_MQTT_HOST = os.getenv("APP_MQTT_HOST")
+APP_MQTT_PORT = os.getenv("APP_MQTT_PORT", 1883)
+APP_MQTT_PROTO = os.getenv("APP_MQTT_PROTO", "mqtt")
+APP_MQTT_USER = os.getenv("APP_MQTT_USER")
+APP_MQTT_PASSWORD = os.getenv("APP_MQTT_PASSWORD")
+
+@app.task
+def pub(details: str) -> None:
+  platform = json.loads(details)['platform']
+  topic = "%s" % (platform)
+  payload = details
+
+  MQTTPublisher().publish(
+      payload = payload,
+      topic = topic
+  )
+
+
+@dataclass
+class MQTTPublisher:
+    """
+    Publishes messages to MQTT
+
+    Attributes:
+        host: MQTT Host
+        port: MQTT Port
+        proto: MQTT Protocol. Currently supports `mqtt`
+        user: MQTT User
+        password: MQTT Password
+    """
+    host: str = APP_MQTT_HOST
+    port: int = APP_MQTT_PORT
+    proto: str = APP_MQTT_PROTO
+    user: str = APP_MQTT_USER
+    password: str = APP_MQTT_PASSWORD
+
+    def __post_init__(self):
+        print('Publisher init complete.')
+    
+    def publish(self, payload, topic, qos=1): 
+        auth = None
+        if self.user:
+            auth = {
+                "username": self.user,
+                "password": self.password 
+            }
+        publish.single(
+            topic=topic, 
+            payload=payload, 
+            qos=qos,
+            hostname=self.host, 
+            port=self.port,
+            auth=auth
+            )

src/kc3zvd/iot_state/wled/publisher.py

@@ -0,0 +1,28 @@
+from celery import Celery
+import os
+import logging
+
+CELERY_BROKER_URL = os.getenv("CELERY_BROKER_URL")
+CELERY_BACKEND_URL = os.getenv("CELERY_BACKEND_URL")
+CELERY_RESULTS_URL = os.getenv("CELERY_RESULTS_URL")
+CELERY_LOG_LEVEL = os.getenv("CELERY_LOG_LEVEL", logging.INFO)
+
+app = Celery('iot-state')
+
+@app.task
+def state_changed() -> None:
+  pass
+
+def mqtt_notify( topic_prefix: str) -> None:
+  # {topic_prefix}/{type}/{area}/{name}/{subtype}
+  #
+  # {type}: One of sensor, switch, notification, meter
+  #
+  # TODO: Ensure inputs are sanitized
+  pass
+
+def discord_notify() -> None:
+  pass
+
+def webhook_notify() -> None:
+  pass

src/kc3zvd/iot_state/workers/__init__.py

@@ -0,0 +1,70 @@
+import logging
+import requests
+from celery import Celery
+import os
+from iot_state.devices import Device
+from mongoengine import connect, MultipleObjectsReturned, DoesNotExist
+
+
+
+CELERY_BROKER_URL = os.getenv("CELERY_BROKER_URL")
+CELERY_BACKEND_URL = os.getenv("CELERY_BACKEND_URL")
+CELERY_RESULTS_URL = os.getenv("CELERY_RESULTS_URL")
+CELERY_LOG_LEVEL = os.getenv("CELERY_LOG_LEVEL", logging.INFO)
+MONGODB_URL = os.getenv("MONGODB_URL")
+
+app = Celery('iot-state')
+
+# Set up logging
+logger = logging.getLogger(__name__)
+logger.setLevel(level=CELERY_LOG_LEVEL)
+
+@app.task
+def discover(details: dict) -> None:
+  
+  for address in details['attributes']['addresses']:
+    url = "http://%s/json/info" % address
+    try:
+      logger.info("Attempting to retrieve info for wled device at %s" % (address))
+      r = requests.get(url)
+      if r.status_code != 200:
+        error_message = "Recieved status code %s from %s, expected 200" % (r.status_code, url)
+        logger.error(error_message)
+        logger.error("Full response: \n %s" % r.content.decode())
+      else:
+        process_discovered_device.delay(details, r.json())
+    except Exception as e:
+      error_message = "Exception while retrieving info from %s." % (url)
+      logger.error(error_message, exc_info=e)
+      continue
+
+@app.task
+def process_discovered_device(details: dict, state: dict) -> None:
+  logger.debug("Details: %s" % (details))
+  logger.debug("State: %s" % (state))
+  connect(host=MONGODB_URL)
+  device_id = state['mac']
+  platform = details['platform']
+  logger.info("Determining if device %s exists in state" % (device_id))
+
+  try:
+    device = Device.objects.get(platform_id = device_id)
+    logger.info("Device %s exists" % (device.platform_id))
+  except MultipleObjectsReturned:
+    logger.warning("Multiple matching devices found")
+  except DoesNotExist:
+    logger.info("Existing device not found, proceeding to bootstrap state")
+    device = Device(
+      platform_id=device_id, 
+      platform=details['platform'],
+      discovery_source=details['source'])
+    device = device.save()
+    logger.info("Device saved with ID %s" % (device_id))
+    register_state_watcher.delay(device_id)
+
+
+@app.task
+def register_state_watcher(device_id):
+  connect(host=MONGODB_URL)
+  device = Device.objects.get(platform_id = device_id)
+  logger.info("Received request to monitor state changes for %s on platform %s" % (device.platform_id, device.platform))