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()