iflow-mcp_enuno-unifi-mcp-server 0.2.1__py3-none-any.whl

src/tools/device_control.py

@@ -0,0 +1,325 @@
+"""Device control MCP tools."""
+
+from typing import Any
+
+from ..api import UniFiClient
+from ..config import Settings
+from ..utils import (
+    ResourceNotFoundError,
+    get_logger,
+    log_audit,
+    sanitize_log_message,
+    validate_confirmation,
+    validate_mac_address,
+    validate_site_id,
+)
+
+
+async def restart_device(
+    site_id: str,
+    device_mac: str,
+    settings: Settings,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Restart a UniFi device.
+
+    Args:
+        site_id: Site identifier
+        device_mac: Device MAC address
+        settings: Application settings
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't restart the device
+
+    Returns:
+        Restart result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+        ResourceNotFoundError: If device not found
+    """
+    site_id = validate_site_id(site_id)
+    device_mac = validate_mac_address(device_mac)
+    validate_confirmation(confirm, "device control operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    parameters = {"site_id": site_id, "device_mac": device_mac}
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(
+                f"DRY RUN: Would restart device '{device_mac}' in site '{site_id}'"
+            )
+        )
+        log_audit(
+            operation="restart_device",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {"dry_run": True, "would_restart": device_mac}
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Verify device exists
+            response = await client.get(f"/ea/sites/{site_id}/devices")
+            # Client now auto-unwraps the "data" field, so response is the actual data
+            devices_data: list[dict[str, Any]] = (
+                response if isinstance(response, list) else response.get("data", [])
+            )
+
+            device_exists = any(
+                validate_mac_address(d.get("mac", "")) == device_mac for d in devices_data
+            )
+            if not device_exists:
+                raise ResourceNotFoundError("device", device_mac)
+
+            # Restart the device
+            restart_data = {"mac": device_mac, "cmd": "restart"}
+            response = await client.post(f"/ea/sites/{site_id}/cmd/devmgr", json_data=restart_data)
+
+            logger.info(
+                sanitize_log_message(
+                    f"Initiated restart for device '{device_mac}' in site '{site_id}'"
+                )
+            )
+            log_audit(
+                operation="restart_device",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "device_mac": device_mac,
+                "message": "Device restart initiated",
+            }
+
+    except Exception as e:
+        logger.error(sanitize_log_message(f"Failed to restart device '{device_mac}': {e}"))
+        log_audit(
+            operation="restart_device",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise
+
+
+async def locate_device(
+    site_id: str,
+    device_mac: str,
+    settings: Settings,
+    enabled: bool = True,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Enable or disable LED locate mode on a device.
+
+    Args:
+        site_id: Site identifier
+        device_mac: Device MAC address
+        settings: Application settings
+        enabled: Enable (True) or disable (False) locate mode
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't change locate state
+
+    Returns:
+        Locate result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+        ResourceNotFoundError: If device not found
+    """
+    site_id = validate_site_id(site_id)
+    device_mac = validate_mac_address(device_mac)
+    validate_confirmation(confirm, "device control operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    parameters = {"site_id": site_id, "device_mac": device_mac, "enabled": enabled}
+
+    action = "enable" if enabled else "disable"
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(
+                f"DRY RUN: Would {action} locate mode for device '{device_mac}' "
+                f"in site '{site_id}'"
+            )
+        )
+        log_audit(
+            operation="locate_device",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {"dry_run": True, f"would_{action}": device_mac}
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Verify device exists
+            response = await client.get(f"/ea/sites/{site_id}/devices")
+            # Client now auto-unwraps the "data" field, so response is the actual data
+            devices_data: list[dict[str, Any]] = (
+                response if isinstance(response, list) else response.get("data", [])
+            )
+
+            device_exists = any(
+                validate_mac_address(d.get("mac", "")) == device_mac for d in devices_data
+            )
+            if not device_exists:
+                raise ResourceNotFoundError("device", device_mac)
+
+            # Set locate state
+            cmd = "set-locate" if enabled else "unset-locate"
+            locate_data = {"mac": device_mac, "cmd": cmd}
+            response = await client.post(f"/ea/sites/{site_id}/cmd/devmgr", json_data=locate_data)
+
+            logger.info(
+                sanitize_log_message(
+                    f"{action.capitalize()}d locate mode for device '{device_mac}' "
+                    f"in site '{site_id}'"
+                )
+            )
+            log_audit(
+                operation="locate_device",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "device_mac": device_mac,
+                "locate_enabled": enabled,
+                "message": f"Locate mode {action}d",
+            }
+
+    except Exception as e:
+        logger.error(
+            sanitize_log_message(f"Failed to {action} locate for device '{device_mac}': {e}")
+        )
+        log_audit(
+            operation="locate_device",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise
+
+
+async def upgrade_device(
+    site_id: str,
+    device_mac: str,
+    settings: Settings,
+    firmware_url: str | None = None,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Trigger firmware upgrade for a device.
+
+    Args:
+        site_id: Site identifier
+        device_mac: Device MAC address
+        settings: Application settings
+        firmware_url: Optional custom firmware URL (uses latest if not provided)
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't initiate upgrade
+
+    Returns:
+        Upgrade result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+        ResourceNotFoundError: If device not found
+    """
+    site_id = validate_site_id(site_id)
+    device_mac = validate_mac_address(device_mac)
+    validate_confirmation(confirm, "device control operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    parameters = {
+        "site_id": site_id,
+        "device_mac": device_mac,
+        "firmware_url": firmware_url,
+    }
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(
+                f"DRY RUN: Would initiate firmware upgrade for device '{device_mac}' "
+                f"in site '{site_id}'"
+            )
+        )
+        log_audit(
+            operation="upgrade_device",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {"dry_run": True, "would_upgrade": device_mac}
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Verify device exists and get details
+            response = await client.get(f"/ea/sites/{site_id}/devices")
+            # Client now auto-unwraps the "data" field, so response is the actual data
+            devices_data: list[dict[str, Any]] = (
+                response if isinstance(response, list) else response.get("data", [])
+            )
+
+            device = None
+            for d in devices_data:
+                if validate_mac_address(d.get("mac", "")) == device_mac:
+                    device = d
+                    break
+
+            if not device:
+                raise ResourceNotFoundError("device", device_mac)
+
+            # Build upgrade command
+            upgrade_data = {"mac": device_mac, "cmd": "upgrade"}
+
+            if firmware_url:
+                upgrade_data["url"] = firmware_url
+
+            response = await client.post(f"/ea/sites/{site_id}/cmd/devmgr", json_data=upgrade_data)
+
+            logger.info(
+                sanitize_log_message(
+                    f"Initiated firmware upgrade for device '{device_mac}' " f"in site '{site_id}'"
+                )
+            )
+            log_audit(
+                operation="upgrade_device",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "device_mac": device_mac,
+                "message": "Firmware upgrade initiated",
+                "current_version": device.get("version"),
+            }
+
+    except Exception as e:
+        logger.error(sanitize_log_message(f"Failed to upgrade device '{device_mac}': {e}"))
+        log_audit(
+            operation="upgrade_device",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise

src/tools/devices.py

@@ -0,0 +1,354 @@
+"""Device management MCP tools."""
+
+from typing import Any
+
+from ..api import UniFiClient
+from ..config import Settings
+from ..models import Device
+from ..utils import (
+    ResourceNotFoundError,
+    audit_action,
+    get_logger,
+    validate_confirmation,
+    validate_device_id,
+    validate_limit_offset,
+    validate_site_id,
+)
+
+
+async def get_device_details(site_id: str, device_id: str, settings: Settings) -> dict[str, Any]:
+    """Get detailed information for a specific device.
+
+    Args:
+        site_id: Site identifier
+        device_id: Device identifier
+        settings: Application settings
+
+    Returns:
+        Device details dictionary
+
+    Raises:
+        ResourceNotFoundError: If device not found
+    """
+    site_id = validate_site_id(site_id)
+    device_id = validate_device_id(device_id)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        # Get all devices and find the specific one
+        response = await client.get(f"/ea/sites/{site_id}/devices")
+        devices_data = response.get("data", []) if isinstance(response, dict) else response
+
+        for device_data in devices_data:
+            if device_data.get("_id") == device_id:
+                device = Device(**device_data)
+                logger.info(f"Retrieved device details for {device_id}")
+                return device.model_dump()  # type: ignore[no-any-return]
+
+        raise ResourceNotFoundError("device", device_id)
+
+
+async def get_device_statistics(site_id: str, device_id: str, settings: Settings) -> dict[str, Any]:
+    """Retrieve real-time statistics for a device.
+
+    Args:
+        site_id: Site identifier
+        device_id: Device identifier
+        settings: Application settings
+
+    Returns:
+        Device statistics dictionary
+    """
+    site_id = validate_site_id(site_id)
+    device_id = validate_device_id(device_id)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        response = await client.get(f"/ea/sites/{site_id}/devices")
+        devices_data = response.get("data", []) if isinstance(response, dict) else response
+
+        for device_data in devices_data:
+            if device_data.get("_id") == device_id:
+                # Extract statistics
+                stats = {
+                    "device_id": device_id,
+                    "uptime": device_data.get("uptime", 0),
+                    "cpu": device_data.get("cpu"),
+                    "mem": device_data.get("mem"),
+                    "tx_bytes": device_data.get("tx_bytes", 0),
+                    "rx_bytes": device_data.get("rx_bytes", 0),
+                    "bytes": device_data.get("bytes", 0),
+                    "state": device_data.get("state"),
+                    "uplink_depth": device_data.get("uplink_depth"),
+                }
+                logger.info(f"Retrieved statistics for device {device_id}")
+                return stats
+
+        raise ResourceNotFoundError("device", device_id)
+
+
+async def list_devices_by_type(
+    site_id: str,
+    device_type: str,
+    settings: Settings,
+    limit: int | None = None,
+    offset: int | None = None,
+) -> list[dict[str, Any]]:
+    """Filter devices by type (AP, switch, gateway).
+
+    Args:
+        site_id: Site identifier
+        device_type: Device type filter (uap, usw, ugw, etc.)
+        settings: Application settings
+        limit: Maximum number of devices to return
+        offset: Number of devices to skip
+
+    Returns:
+        List of device dictionaries
+    """
+    site_id = validate_site_id(site_id)
+    limit, offset = validate_limit_offset(limit, offset)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        response = await client.get(f"/ea/sites/{site_id}/devices")
+        devices_data = response.get("data", []) if isinstance(response, dict) else response
+
+        # Filter by type
+        filtered = [
+            d
+            for d in devices_data
+            if d.get("type", "").lower() == device_type.lower()
+            or device_type.lower() in d.get("model", "").lower()
+        ]
+
+        # Apply pagination
+        paginated = filtered[offset : offset + limit]
+
+        # Parse into Device models
+        devices = [Device(**d).model_dump() for d in paginated]
+
+        logger.info(
+            f"Retrieved {len(devices)} devices of type '{device_type}' " f"for site '{site_id}'"
+        )
+        return devices
+
+
+async def search_devices(
+    site_id: str,
+    query: str,
+    settings: Settings,
+    limit: int | None = None,
+    offset: int | None = None,
+) -> list[dict[str, Any]]:
+    """Search devices by name, MAC, or IP address.
+
+    Args:
+        site_id: Site identifier
+        query: Search query string
+        settings: Application settings
+        limit: Maximum number of devices to return
+        offset: Number of devices to skip
+
+    Returns:
+        List of matching device dictionaries
+    """
+    site_id = validate_site_id(site_id)
+    limit, offset = validate_limit_offset(limit, offset)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        response = await client.get(f"/ea/sites/{site_id}/devices")
+        devices_data = response.get("data", []) if isinstance(response, dict) else response
+
+        # Search by name, MAC, or IP
+        query_lower = query.lower()
+        filtered = [
+            d
+            for d in devices_data
+            if query_lower in d.get("name", "").lower()
+            or query_lower in d.get("mac", "").lower()
+            or query_lower in d.get("ip", "").lower()
+            or query_lower in d.get("model", "").lower()
+        ]
+
+        # Apply pagination
+        paginated = filtered[offset : offset + limit]
+
+        # Parse into Device models
+        devices = [Device(**d).model_dump() for d in paginated]
+
+        logger.info(f"Found {len(devices)} devices matching '{query}' in site '{site_id}'")
+        return devices
+
+
+async def list_pending_devices(
+    site_id: str,
+    settings: Settings,
+    limit: int | None = None,
+    offset: int | None = None,
+) -> list[dict[str, Any]]:
+    """List devices awaiting adoption on the specified site.
+
+    Args:
+        site_id: Site identifier
+        settings: Application settings
+        limit: Maximum number of devices to return
+        offset: Number of devices to skip
+
+    Returns:
+        List of pending device dictionaries
+    """
+    site_id = validate_site_id(site_id)
+    limit, offset = validate_limit_offset(limit, offset)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        params = {}
+        if limit is not None:
+            params["limit"] = limit
+        if offset is not None:
+            params["offset"] = offset
+
+        response = await client.get(
+            f"/integration/v1/sites/{site_id}/devices/pending", params=params
+        )
+        devices_data = response.get("data", [])
+
+        # Parse into Device models
+        devices = [Device(**d).model_dump() for d in devices_data]
+
+        logger.info(f"Retrieved {len(devices)} pending devices for site '{site_id}'")
+        return devices
+
+
+async def adopt_device(
+    site_id: str,
+    device_id: str,
+    settings: Settings,
+    name: str | None = None,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Adopt a pending device onto the specified site.
+
+    Args:
+        site_id: Site identifier
+        device_id: Device identifier to adopt
+        settings: Application settings
+        name: Optional device name
+        confirm: Confirmation flag (required)
+        dry_run: If True, validate but don't execute
+
+    Returns:
+        Adopted device information
+    """
+    validate_confirmation(confirm, "adopt device")
+    site_id = validate_site_id(site_id)
+    device_id = validate_device_id(device_id)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        payload = {}
+        if name:
+            payload["name"] = name
+
+        if dry_run:
+            logger.info(f"[DRY RUN] Would adopt device {device_id} with payload: {payload}")
+            return {"dry_run": True, "device_id": device_id, "payload": payload}
+
+        response = await client.post(
+            f"/integration/v1/sites/{site_id}/devices/{device_id}/adopt", json_data=payload
+        )
+        data = response.get("data", response)
+
+        # Audit the action
+        await audit_action(
+            settings,
+            action_type="adopt_device",
+            resource_type="device",
+            resource_id=device_id,
+            site_id=site_id,
+            details={"name": name} if name else {},
+        )
+
+        logger.info(f"Successfully adopted device {device_id}")
+        return Device(**data).model_dump()  # type: ignore[no-any-return]
+
+
+async def execute_port_action(
+    site_id: str,
+    device_id: str,
+    port_idx: int,
+    action: str,
+    settings: Settings,
+    params: dict[str, Any] | None = None,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Execute an action on a specific port of a device.
+
+    Args:
+        site_id: Site identifier
+        device_id: Device identifier
+        port_idx: Port index number
+        action: Action to perform (power-cycle, enable, disable)
+        settings: Application settings
+        params: Additional action parameters
+        confirm: Confirmation flag (required)
+        dry_run: If True, validate but don't execute
+
+    Returns:
+        Action result
+    """
+    validate_confirmation(confirm, f"execute port action '{action}'")
+    site_id = validate_site_id(site_id)
+    device_id = validate_device_id(device_id)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        payload = {"action": action, "params": params or {}}
+
+        if dry_run:
+            logger.info(
+                f"[DRY RUN] Would execute port action '{action}' on device {device_id} port {port_idx}"
+            )
+            return {
+                "dry_run": True,
+                "device_id": device_id,
+                "port_idx": port_idx,
+                "payload": payload,
+            }
+
+        response = await client.post(
+            f"/integration/v1/sites/{site_id}/devices/{device_id}/ports/{port_idx}/action",
+            json_data=payload,
+        )
+        data = response.get("data", response)
+
+        # Audit the action
+        await audit_action(
+            settings,
+            action_type="port_action",
+            resource_type="device_port",
+            resource_id=f"{device_id}:{port_idx}",
+            site_id=site_id,
+            details={"action": action},
+        )
+
+        logger.info(f"Successfully executed port action '{action}' on port {port_idx}")
+        return {"success": True, "action": action, "port_idx": port_idx, "result": data}