TasmotaManager/unifi_client.py

"""UniFi Controller API client."""

import requests
import urllib3
import logging
from typing import List, Dict, Optional

# Disable SSL warnings for self-signed certificates
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class AuthenticationError(Exception):
    """Raised when authentication with UniFi controller fails."""
    pass


class UniFiDataError(Exception):
    """Raised when UniFi controller returns unexpected data."""
    pass


class UnifiClient:
    """Client for interacting with UniFi Controller API."""
    
    def __init__(self, host: str, username: str, password: str, site: str = 'default', 
                 verify_ssl: bool = False, logger: Optional[logging.Logger] = None):
        """
        Initialize UniFi client.
        
        Args:
            host: UniFi controller URL (e.g., 'https://192.168.1.1')
            username: Username for authentication
            password: Password for authentication
            site: Site name (default: 'default')
            verify_ssl: Whether to verify SSL certificates
            logger: Optional logger instance
        """
        self.base_url = host.rstrip('/')
        self.username = username
        self.password = password
        self.site_id = site
        self.verify_ssl = verify_ssl
        self.token = None
        self.session = requests.Session()
        self.logger = logger or logging.getLogger(__name__)
        
        # Login to get session token
        self._login()
    
    def _request_json(self, endpoint: str, method: str = 'GET', 
                     data: Optional[dict] = None) -> dict:
        """
        Make a request to the UniFi API and return JSON response.
        
        Args:
            endpoint: API endpoint path
            method: HTTP method (GET, POST, etc.)
            data: Optional data for POST requests
            
        Returns:
            dict: JSON response
            
        Raises:
            UniFiDataError: If request fails or returns invalid data
        """
        url = f"{self.base_url}{endpoint}"
        
        try:
            if method == 'GET':
                response = self.session.get(url, verify=self.verify_ssl, timeout=30)
            elif method == 'POST':
                response = self.session.post(url, json=data, verify=self.verify_ssl, timeout=30)
            else:
                raise ValueError(f"Unsupported HTTP method: {method}")
            
            response.raise_for_status()
            
            try:
                json_response = response.json()
            except ValueError:
                raise UniFiDataError(f"Invalid JSON response from {endpoint}")
            
            # Check for UniFi API error response
            if isinstance(json_response, dict):
                if json_response.get('meta', {}).get('rc') != 'ok':
                    error_msg = json_response.get('meta', {}).get('msg', 'Unknown error')
                    raise UniFiDataError(f"UniFi API error: {error_msg}")
            
            return json_response
            
        except requests.exceptions.RequestException as e:
            self.logger.error(f"Request to {endpoint} failed: {e}")
            raise UniFiDataError(f"Request failed: {e}")
    
    def _login(self):
        """
        Authenticate with the UniFi controller.
        
        Raises:
            AuthenticationError: If authentication fails
        """
        login_data = {
            'username': self.username,
            'password': self.password
        }
        
        try:
            response = self._request_json('/api/auth/login', method='POST', data=login_data)
            self.logger.debug("Successfully authenticated with UniFi controller")
            
        except UniFiDataError as e:
            self.logger.error(f"Authentication failed: {e}")
            raise AuthenticationError(f"Failed to authenticate: {e}")
    
    def get_clients(self) -> List[Dict]:
        """
        Get all clients from the UniFi controller.
        
        Returns:
            list: List of client dictionaries
            
        Raises:
            UniFiDataError: If request fails
        """
        endpoint = f'/api/s/{self.site_id}/stat/sta'
        
        try:
            response = self._request_json(endpoint)
            
            if isinstance(response, dict) and 'data' in response:
                clients = response['data']
                self.logger.debug(f"Retrieved {len(clients)} clients from UniFi controller")
                return clients
            else:
                raise UniFiDataError("Unexpected response format from UniFi controller")
                
        except UniFiDataError as e:
            self.logger.error(f"Failed to get clients: {e}")
            raise
    
    def get_devices(self) -> List[Dict]:
        """
        Get all devices (APs, switches, etc.) from the UniFi controller.
        
        Returns:
            list: List of device dictionaries
            
        Raises:
            UniFiDataError: If request fails
        """
        endpoint = f'/api/s/{self.site_id}/stat/device'
        
        try:
            response = self._request_json(endpoint)
            
            if isinstance(response, dict) and 'data' in response:
                devices = response['data']
                self.logger.debug(f"Retrieved {len(devices)} devices from UniFi controller")
                return devices
            else:
                raise UniFiDataError("Unexpected response format from UniFi controller")
                
        except UniFiDataError as e:
            self.logger.error(f"Failed to get devices: {e}")
            raise