system2mqtt/system2mqtt.py

#!/usr/bin/env python3

import json
import os
import socket
import sys
import time
import yaml
import platform
import paho.mqtt.client as mqtt
from typing import Dict, Any, List
import importlib.util
import glob
from datetime import datetime, timedelta

class System2MQTT:
    def __init__(self, config_path: str = "config.yaml"):
        self.config = self._load_config(config_path)
        self.hostname = socket.gethostname()
        self.client = self._setup_mqtt_client()
        self.connected = False
        self.device_info = self._get_device_info()
        self.collectors = self._load_collectors()
        self.last_run = {}  # Speichert den Zeitpunkt des letzten Laufs für jeden Sammler

    def _load_config(self, config_path: str) -> Dict[str, Any]:
        """Load configuration from YAML file."""
        try:
            with open(config_path, 'r') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Error loading config: {e}")
            sys.exit(1)

    def _setup_mqtt_client(self) -> mqtt.Client:
        """Setup MQTT client with configuration."""
        client = mqtt.Client(
            client_id=self.config['mqtt']['client_id'].format(hostname=self.hostname)
        )
        client.username_pw_set(
            self.config['mqtt']['username'],
            self.config['mqtt']['password']
        )
        client.on_connect = self._on_connect
        client.on_disconnect = self._on_disconnect
        return client

    def _on_connect(self, client, userdata, flags, rc):
        """Callback for when the client connects to the broker."""
        if rc == 0:
            print("Connected to MQTT broker")
            self.connected = True
        else:
            print(f"Failed to connect to MQTT broker with code: {rc}")
            self.connected = False

    def _on_disconnect(self, client, userdata, rc):
        """Callback for when the client disconnects from the broker."""
        print(f"Disconnected from MQTT broker with code: {rc}")
        self.connected = False

    def _get_device_info(self) -> Dict[str, Any]:
        """Get device information for Home Assistant."""
        return {
            "identifiers": [f"system2mqtt_{self.hostname}"],
            "name": f"System {self.hostname}",
            "model": platform.machine(),
            "manufacturer": platform.system()
        }

    def _get_unique_id(self, sensor_id: str) -> str:
        """Generate unique_id from sensor_id."""
        return f"system2mqtt_{self.hostname}_{sensor_id}"

    def _get_state_topic(self, sensor_id: str) -> str:
        """Generate state topic from sensor_id."""
        return f"system2mqtt/{self.hostname}/{sensor_id}/state"

    def _load_collectors(self) -> List[Dict[str, Any]]:
        """Load all collector modules from the collectors directory."""
        collectors = []
        collector_dir = os.path.join(os.path.dirname(__file__), 'collectors')
        
        # Find all Python files in the collectors directory
        collector_files = glob.glob(os.path.join(collector_dir, '*.py'))
        
        for collector_file in collector_files:
            if collector_file.endswith('__init__.py'):
                continue
                
            module_name = os.path.splitext(os.path.basename(collector_file))[0]
            spec = importlib.util.spec_from_file_location(module_name, collector_file)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            
            if hasattr(module, 'collect_metrics'):
                # Get interval from config or use collector's default
                default_interval = getattr(module, 'DEFAULT_INTERVAL', self.config['collectors']['default_interval'])
                interval = self.config['collectors']['intervals'].get(
                    module_name,
                    default_interval
                )
                
                collectors.append({
                    'module': module,
                    'name': module_name,
                    'interval': interval
                })
                print(f"Loaded collector: {module_name} (interval: {interval}s)")
        
        return collectors

    def process_collector_data(self, data: Dict[str, Any]):
        """Process data from collectors and publish to MQTT."""
        if not self.connected:
            print("Not connected to MQTT broker")
            return

        # Publish discovery messages for each entity
        for entity in data['entities']:
            sensor_id = entity['sensor_id']
            unique_id = self._get_unique_id(sensor_id)
            state_topic = self._get_state_topic(sensor_id)
            discovery_topic = f"{self.config['mqtt']['discovery_prefix']}/sensor/{unique_id}/config"
            
            # Prepare discovery message
            discovery_msg = {
                "name": entity['name'],
                "unique_id": unique_id,
                "state_topic": state_topic,
                "state_class": entity['state_class'],
                "unit_of_measurement": entity['unit_of_measurement'],
                "device_class": entity['device_class'],
                "device": self.device_info
            }

            # Publish discovery message
            self.client.publish(
                discovery_topic,
                json.dumps(discovery_msg),
                retain=True
            )

            # Publish state
            self.client.publish(
                state_topic,
                entity['value'],
                retain=True
            )

            # Publish attributes if present
            if 'attributes' in entity:
                attributes_topic = f"{state_topic}/attributes"
                self.client.publish(
                    attributes_topic,
                    json.dumps(entity['attributes']),
                    retain=True
                )

    def should_run_collector(self, collector: Dict[str, Any]) -> bool:
        """Check if a collector should run based on its interval."""
        now = datetime.now()
        last_run = self.last_run.get(collector['name'])
        
        if last_run is None:
            return True
            
        interval = timedelta(seconds=collector['interval'])
        return (now - last_run) >= interval

    def collect_and_publish(self):
        """Collect metrics from all collectors and publish them."""
        for collector in self.collectors:
            if not self.should_run_collector(collector):
                continue
                
            try:
                data = collector['module'].collect_metrics()
                self.process_collector_data(data)
                self.last_run[collector['name']] = datetime.now()
                print(f"Updated {collector['name']} metrics")
            except Exception as e:
                print(f"Error collecting metrics from {collector['name']}: {e}")

    def connect(self):
        """Connect to MQTT broker."""
        try:
            self.client.connect(
                self.config['mqtt']['broker'],
                self.config['mqtt']['port']
            )
            self.client.loop_start()
        except Exception as e:
            print(f"Error connecting to MQTT broker: {e}")
            sys.exit(1)

    def disconnect(self):
        """Disconnect from MQTT broker."""
        self.client.loop_stop()
        self.client.disconnect()

def main():
    """Main function."""
    system2mqtt = System2MQTT()
    system2mqtt.connect()

    try:
        # Initial collection
        system2mqtt.collect_and_publish()
        
        # Main loop - check every second if any collector needs to run
        while True:
            system2mqtt.collect_and_publish()
            time.sleep(1)  # Check every second
            
    except KeyboardInterrupt:
        print("\nShutting down...")
    finally:
        system2mqtt.disconnect()

if __name__ == "__main__":
    main()