system2mqtt/main.py

#!/usr/bin/env python3

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

# Load configuration
with open('config.yaml', 'r') as f:
    config = yaml.safe_load(f)

# MQTT Configuration
MQTT_HOST = config['mqtt']['host']
MQTT_PORT = config['mqtt']['port']
MQTT_USERNAME = config['mqtt']['username']
MQTT_PASSWORD = config['mqtt']['password']
MQTT_CLIENT_ID = config['mqtt']['client_id']
MQTT_DISCOVERY_PREFIX = config['mqtt']['discovery_prefix']

# Get hostname
HOSTNAME = socket.gethostname()

# Get OS information
OS_INFO = {
    'system': platform.system(),
    'release': platform.release(),
    'version': platform.version(),
    'machine': platform.machine(),
    'processor': platform.processor()
}

def get_collectors() -> List[str]:
    """Get list of available collector scripts."""
    collectors = []
    for file in os.listdir('collectors'):
        if file.endswith('.py') and not file.startswith('__'):
            collectors.append(file[:-3])
    return collectors

def load_collector(collector_name: str) -> Any:
    """Load a collector module."""
    try:
        return importlib.import_module(f'collectors.{collector_name}')
    except ImportError as e:
        print(f"Error loading collector {collector_name}: {e}")
        return None

def get_device_info() -> Dict[str, Any]:
    """Get device information for Home Assistant."""
    return {
        "identifiers": [f"system2mqtt_{HOSTNAME}"],
        "name": f"System Metrics - {HOSTNAME}",
        "model": OS_INFO['system'],
        "manufacturer": "system2mqtt",
        "sw_version": OS_INFO['version'],
        "configuration_url": f"http://{HOSTNAME}",
        "hw_version": OS_INFO['machine']
    }

def get_device_attributes() -> Dict[str, Any]:
    """Get additional device attributes."""
    return {
        "operating_system": OS_INFO['system'],
        "os_release": OS_INFO['release'],
        "os_version": OS_INFO['version'],
        "architecture": OS_INFO['machine'],
        "processor": OS_INFO['processor']
    }

def on_connect(client: mqtt.Client, userdata: Any, flags: Dict[str, Any], rc: int) -> None:
    """Callback for when the client connects to the MQTT broker."""
    print(f"Connected with result code {rc}")
    # Subscribe to any topics if needed
    # client.subscribe("$SYS/#")

def on_disconnect(client: mqtt.Client, userdata: Any, rc: int) -> None:
    """Callback for when the client disconnects from the MQTT broker."""
    print(f"Disconnected with result code {rc}")
    if rc != 0:
        print("Unexpected disconnection. Attempting to reconnect...")

def main():
    """Main function."""
    # Create MQTT client
    client = mqtt.Client(client_id=MQTT_CLIENT_ID)
    client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)
    client.on_connect = on_connect
    client.on_disconnect = on_disconnect

    # Connect to MQTT broker
    try:
        client.connect(MQTT_HOST, MQTT_PORT, 60)
    except Exception as e:
        print(f"Failed to connect to MQTT broker: {e}")
        sys.exit(1)

    # Start the MQTT client loop
    client.loop_start()

    # Get list of collectors
    collectors = get_collectors()
    print(f"Found collectors: {collectors}")

    # Main loop
    while True:
        try:
            # Load and run each collector
            for collector_name in collectors:
                collector = load_collector(collector_name)
                if collector:
                    try:
                        # Get metrics from collector
                        metrics = collector.collect_metrics()
                        
                        # Add device info to each entity
                        device_info = get_device_info()
                        device_attributes = get_device_attributes()
                        
                        for entity in metrics.get('entities', []):
                            # Create discovery topic
                            discovery_topic = f"{MQTT_DISCOVERY_PREFIX}/sensor/system2mqtt_{HOSTNAME}_{entity['sensor_id']}/config"
                            
                            # Create discovery payload
                            discovery_payload = {
                                "name": entity['name'],
                                "unique_id": f"system2mqtt_{HOSTNAME}_{entity['sensor_id']}",
                                "state_topic": f"system2mqtt/{HOSTNAME}/{entity['sensor_id']}/state",
                                "state_class": entity['state_class'],
                                "unit_of_measurement": entity['unit_of_measurement'],
                                "device_class": entity['device_class'],
                                "icon": entity.get('icon'),
                                "device": device_info,
                                "availability": {
                                    "topic": f"system2mqtt/{HOSTNAME}/status",
                                    "payload_available": "online",
                                    "payload_not_available": "offline"
                                },
                                "json_attributes_topic": f"system2mqtt/{HOSTNAME}/{entity['sensor_id']}/attributes"
                            }
                            
                            # Publish discovery message
                            client.publish(discovery_topic, json.dumps(discovery_payload), retain=True)
                            
                            # Publish state
                            state_topic = f"system2mqtt/{HOSTNAME}/{entity['sensor_id']}/state"
                            client.publish(state_topic, entity['value'], retain=True)
                            
                            # Publish attributes
                            attributes = entity.get('attributes', {})
                            attributes.update(device_attributes)  # Add device attributes
                            attributes_topic = f"system2mqtt/{HOSTNAME}/{entity['sensor_id']}/attributes"
                            client.publish(attributes_topic, json.dumps(attributes), retain=True)
                            
                    except Exception as e:
                        print(f"Error collecting metrics from {collector_name}: {e}")
                        continue

            # Publish availability status
            client.publish(f"system2mqtt/{HOSTNAME}/status", "online", retain=True)

            # Wait before next collection
            time.sleep(60)  # Collect metrics every minute

        except KeyboardInterrupt:
            print("Stopping...")
            # Publish offline status before stopping
            client.publish(f"system2mqtt/{HOSTNAME}/status", "offline", retain=True)
            break
        except Exception as e:
            print(f"Error in main loop: {e}")
            time.sleep(60)  # Wait before retrying

    # Stop the MQTT client loop and disconnect
    client.loop_stop()
    client.disconnect()

if __name__ == "__main__":
    main()