Initial Creation

TasmotaManager.py

@@ -0,0 +1,185 @@
+import json
+import logging
+import os
+import sys
+from datetime import datetime
+from typing import Optional
+import requests
+from urllib3.exceptions import InsecureRequestWarning
+requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)
+
+class UnifiClient:
+    def __init__(self, host: str, username: str, password: str, site_id: str = 'default', ssl_verify: bool = True):
+        self.base_url = f"https://{host}"
+        self.site_id = site_id
+        self.session = requests.Session()
+        self.session.verify = ssl_verify
+        self._login(username, password)
+
+    def _login(self, username: str, password: str) -> None:
+        """Authenticate with the UniFi Controller."""
+        login_url = f"{self.base_url}/api/login"
+        response = self.session.post(
+            login_url,
+            json={"username": username, "password": password}
+        )
+        response.raise_for_status()
+
+    def get_clients(self) -> list:
+        """Get all clients from the UniFi Controller."""
+        url = f"{self.base_url}/api/s/{self.site_id}/stat/sta"
+        response = self.session.get(url)
+        response.raise_for_status()
+        return response.json().get('data', [])
+
+class TasmotaDiscovery:
+    def __init__(self, debug: bool = False):
+        """Initialize the TasmotaDiscovery with optional debug mode."""
+        log_level = logging.DEBUG if debug else logging.INFO
+        logging.basicConfig(
+            level=log_level,
+            format='%(asctime)s - %(levelname)s - %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S'
+        )
+        self.logger = logging.getLogger(__name__)
+        self.config = None
+        self.unifi_client = None
+
+    def load_config(self, config_path: Optional[str] = None) -> dict:
+        """Load configuration from JSON file."""
+        if config_path is None:
+            config_path = os.path.join(os.path.dirname(__file__), 'config.json')
+        
+        self.logger.debug(f"Loading configuration from: {config_path}")
+        try:
+            with open(config_path, 'r') as config_file:
+                self.config = json.load(config_file)
+                self.logger.debug("Configuration loaded successfully from %s", config_path)
+                return self.config
+        except FileNotFoundError:
+            self.logger.error(f"Configuration file not found at {config_path}")
+            sys.exit(1)
+        except json.JSONDecodeError:
+            self.logger.error("Invalid JSON in configuration file")
+            sys.exit(1)
+
+    def setup_unifi_client(self) -> UnifiClient:
+        """Initialize UniFi client with configuration."""
+        self.logger.debug("Setting up UniFi client")
+        try:
+            host = self.config['unifi']['host']
+            # Remove https:// if present
+            if host.startswith('https://'):
+                host = host[8:]
+            elif host.startswith('http://'):
+                host = host[7:]
+
+            self.logger.debug(f"Connecting to UniFi Controller at {host}")
+            
+            self.unifi_client = UnifiClient(
+                host=host,
+                username=self.config['unifi']['username'],
+                password=self.config['unifi']['password'],
+                site_id=self.config['unifi'].get('site', 'default'),
+                ssl_verify=False
+            )
+            self.logger.debug("UniFi client setup successful")
+            return self.unifi_client
+        except KeyError as e:
+            self.logger.error(f"Missing required UniFi configuration: {e}")
+            self.logger.debug("Connection details:")
+            self.logger.debug(f"Host: {self.config['unifi'].get('host', 'Not set')}")
+            self.logger.debug(f"Username: {self.config['unifi'].get('username', 'Not set')}")
+            self.logger.debug("Please verify your configuration file")
+            sys.exit(1)
+        except Exception as e:
+            self.logger.error(f"Error connecting to UniFi controller: {e}")
+            sys.exit(1)
+
+    def is_tasmota_device(self, device: dict) -> bool:
+        """Determine if a device is a Tasmota device."""
+        name = device.get('name', '').lower()
+        hostname = device.get('hostname', '').lower()
+        self.logger.debug(f"Checking device: {name} ({hostname})")
+        matches = any([
+            name.startswith('tasmota'),
+            name.startswith('sonoff'),
+            name.endswith('-ts'),
+            hostname.startswith('tasmota'),
+            hostname.startswith('sonoff')
+        ])
+        if matches:
+            self.logger.debug(f"Found Tasmota device: {name}")
+        return matches
+
+    def get_tasmota_devices(self) -> list:
+        """Query UniFi controller and filter Tasmota devices."""
+        devices = []
+        self.logger.debug("Querying UniFi controller for devices")
+        try:
+            all_clients = self.unifi_client.get_clients()
+            self.logger.debug(f"Found {len(all_clients)} total devices")
+            
+            for device in all_clients:
+                if self.is_tasmota_device(device):
+                    device_info = {
+                        "name": device.get('name', device.get('hostname', 'Unknown')),
+                        "ip": device.get('ip', ''),
+                        "mac": device.get('mac', ''),
+                        "last_seen": device.get('last_seen', ''),
+                        "hostname": device.get('hostname', ''),
+                        "notes": device.get('note', ''),
+                    }
+                    devices.append(device_info)
+            
+            self.logger.debug(f"Found {len(devices)} Tasmota devices")
+            return devices
+        except Exception as e:
+            self.logger.error(f"Error getting devices from UniFi controller: {e}")
+            return []
+
+    def save_tasmota_config(self, devices: list, filename: str = "tasmota_devices.json") -> None:
+        """Save Tasmota device information to a JSON file."""
+        self.logger.debug(f"Saving Tasmota configuration to {filename}")
+        config = {
+            "tasmota": {
+                "devices": devices,
+                "generated_at": datetime.now().isoformat(),
+                "total_devices": len(devices)
+            }
+        }
+        
+        try:
+            if os.path.exists(filename):
+                backup_name = f"{filename}.backup"
+                os.rename(filename, backup_name)
+                self.logger.info(f"Created backup of existing configuration as {backup_name}")
+            
+            with open(filename, 'w') as f:
+                json.dump(config, f, indent=4)
+            self.logger.info(f"Successfully saved {len(devices)} Tasmota devices to {filename}")
+            
+            print("\nFound Tasmota Devices:")
+            for device in devices:
+                print(f"Name: {device['name']:<20} IP: {device['ip']:<15} MAC: {device['mac']}")
+                
+        except Exception as e:
+            self.logger.error(f"Error saving Tasmota configuration: {e}")
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description='Discover Tasmota devices on UniFi network')
+    parser.add_argument('--debug', action='store_true', help='Enable debug logging')
+    parser.add_argument('--config', type=str, help='Path to configuration file')
+    parser.add_argument('--output', type=str, default='tasmota_devices.json',
+                      help='Output file for device list (default: tasmota_devices.json)')
+    args = parser.parse_args()
+
+    discovery = TasmotaDiscovery(debug=args.debug)
+    discovery.load_config(args.config)
+    discovery.setup_unifi_client()
+    devices = discovery.get_tasmota_devices()
+    discovery.save_tasmota_config(devices, args.output)
+
+if __name__ == '__main__':
+    main()