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

src/tools/client_management.py

@@ -0,0 +1,505 @@
+"""Client management 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 block_client(
+    site_id: str,
+    client_mac: str,
+    settings: Settings,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Block a client from accessing the network.
+
+    Args:
+        site_id: Site identifier
+        client_mac: Client MAC address
+        settings: Application settings
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't block the client
+
+    Returns:
+        Block result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+        ResourceNotFoundError: If client not found
+    """
+    site_id = validate_site_id(site_id)
+    client_mac = validate_mac_address(client_mac)
+    validate_confirmation(confirm, "client management operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    parameters = {"site_id": site_id, "client_mac": client_mac}
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(f"DRY RUN: Would block client '{client_mac}' in site '{site_id}'")
+        )
+        log_audit(
+            operation="block_client",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {"dry_run": True, "would_block": client_mac}
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Verify client exists (check both active and all users)
+            response = await client.get(f"/ea/sites/{site_id}/stat/alluser")
+            # Client now auto-unwraps the "data" field, so response is the actual data
+            clients_data: list[dict[str, Any]] = (
+                response if isinstance(response, list) else response.get("data", [])
+            )
+
+            client_exists = any(
+                validate_mac_address(c.get("mac", "")) == client_mac for c in clients_data
+            )
+            if not client_exists:
+                raise ResourceNotFoundError("client", client_mac)
+
+            # Block the client
+            block_data = {"mac": client_mac, "cmd": "block-sta"}
+            response = await client.post(f"/ea/sites/{site_id}/cmd/stamgr", json_data=block_data)
+
+            logger.info(sanitize_log_message(f"Blocked client '{client_mac}' in site '{site_id}'"))
+            log_audit(
+                operation="block_client",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "client_mac": client_mac,
+                "message": "Client blocked from network",
+            }
+
+    except Exception as e:
+        logger.error(sanitize_log_message(f"Failed to block client '{client_mac}': {e}"))
+        log_audit(
+            operation="block_client",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise
+
+
+async def unblock_client(
+    site_id: str,
+    client_mac: str,
+    settings: Settings,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Unblock a previously blocked client.
+
+    Args:
+        site_id: Site identifier
+        client_mac: Client MAC address
+        settings: Application settings
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't unblock the client
+
+    Returns:
+        Unblock result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+    """
+    site_id = validate_site_id(site_id)
+    client_mac = validate_mac_address(client_mac)
+    validate_confirmation(confirm, "client management operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    parameters = {"site_id": site_id, "client_mac": client_mac}
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(
+                f"DRY RUN: Would unblock client '{client_mac}' in site '{site_id}'"
+            )
+        )
+        log_audit(
+            operation="unblock_client",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {"dry_run": True, "would_unblock": client_mac}
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Unblock the client
+            unblock_data = {"mac": client_mac, "cmd": "unblock-sta"}
+            await client.post(f"/ea/sites/{site_id}/cmd/stamgr", json_data=unblock_data)
+
+            logger.info(
+                sanitize_log_message(f"Unblocked client '{client_mac}' in site '{site_id}'")
+            )
+            log_audit(
+                operation="unblock_client",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "client_mac": client_mac,
+                "message": "Client unblocked",
+            }
+
+    except Exception as e:
+        logger.error(sanitize_log_message(f"Failed to unblock client '{client_mac}': {e}"))
+        log_audit(
+            operation="unblock_client",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise
+
+
+async def reconnect_client(
+    site_id: str,
+    client_mac: str,
+    settings: Settings,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Force a client to reconnect (disconnect and re-authenticate).
+
+    Args:
+        site_id: Site identifier
+        client_mac: Client MAC address
+        settings: Application settings
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't force reconnection
+
+    Returns:
+        Reconnect result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+        ResourceNotFoundError: If client not found
+    """
+    site_id = validate_site_id(site_id)
+    client_mac = validate_mac_address(client_mac)
+    validate_confirmation(confirm, "client management operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    parameters = {"site_id": site_id, "client_mac": client_mac}
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(
+                f"DRY RUN: Would force reconnect for client '{client_mac}' in site '{site_id}'"
+            )
+        )
+        log_audit(
+            operation="reconnect_client",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {"dry_run": True, "would_reconnect": client_mac}
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Verify client is currently connected
+            response = await client.get(f"/ea/sites/{site_id}/sta")
+            # Client now auto-unwraps the "data" field, so response is the actual data
+            clients_data: list[dict[str, Any]] = (
+                response if isinstance(response, list) else response.get("data", [])
+            )
+
+            client_exists = any(
+                validate_mac_address(c.get("mac", "")) == client_mac for c in clients_data
+            )
+            if not client_exists:
+                raise ResourceNotFoundError("active client", client_mac)
+
+            # Force client reconnection
+            reconnect_data = {"mac": client_mac, "cmd": "kick-sta"}
+            response = await client.post(
+                f"/ea/sites/{site_id}/cmd/stamgr", json_data=reconnect_data
+            )
+
+            logger.info(
+                sanitize_log_message(
+                    f"Forced reconnect for client '{client_mac}' in site '{site_id}'"
+                )
+            )
+            log_audit(
+                operation="reconnect_client",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "client_mac": client_mac,
+                "message": "Client forced to reconnect",
+            }
+
+    except Exception as e:
+        logger.error(sanitize_log_message(f"Failed to reconnect client '{client_mac}': {e}"))
+        log_audit(
+            operation="reconnect_client",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise
+
+
+async def authorize_guest(
+    site_id: str,
+    client_mac: str,
+    duration: int,
+    settings: Settings,
+    upload_limit_kbps: int | None = None,
+    download_limit_kbps: int | None = None,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Authorize a guest client for network access.
+
+    Args:
+        site_id: Site identifier
+        client_mac: Client MAC address
+        duration: Access duration in seconds
+        settings: Application settings
+        upload_limit_kbps: Upload speed limit in kbps
+        download_limit_kbps: Download speed limit in kbps
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't authorize
+
+    Returns:
+        Authorization result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+    """
+    site_id = validate_site_id(site_id)
+    client_mac = validate_mac_address(client_mac)
+    validate_confirmation(confirm, "client management operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    parameters = {
+        "site_id": site_id,
+        "client_mac": client_mac,
+        "duration": duration,
+    }
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(
+                f"DRY RUN: Would authorize guest client '{client_mac}' for {duration}s in site '{site_id}'"
+            )
+        )
+        log_audit(
+            operation="authorize_guest",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {"dry_run": True, "would_authorize": client_mac, "duration": duration}
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Build authorization payload
+            auth_data = {
+                "action": "authorize-guest",
+                "params": {
+                    "duration": duration,
+                },
+            }
+
+            if upload_limit_kbps is not None:
+                auth_data["params"]["uploadLimit"] = upload_limit_kbps  # type: ignore[index]
+            if download_limit_kbps is not None:
+                auth_data["params"]["downloadLimit"] = download_limit_kbps  # type: ignore[index]
+
+            # Authorize guest using new API endpoint
+            await client.post(
+                f"/integration/v1/sites/{site_id}/clients/{client_mac}/action",
+                json_data=auth_data,
+            )
+
+            logger.info(
+                sanitize_log_message(
+                    f"Authorized guest client '{client_mac}' for {duration}s in site '{site_id}'"
+                )
+            )
+            log_audit(
+                operation="authorize_guest",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "client_mac": client_mac,
+                "duration": duration,
+                "message": f"Guest authorized for {duration} seconds",
+            }
+
+    except Exception as e:
+        logger.error(sanitize_log_message(f"Failed to authorize guest client '{client_mac}': {e}"))
+        log_audit(
+            operation="authorize_guest",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise
+
+
+async def limit_bandwidth(
+    site_id: str,
+    client_mac: str,
+    settings: Settings,
+    upload_limit_kbps: int | None = None,
+    download_limit_kbps: int | None = None,
+    confirm: bool = False,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Apply bandwidth restrictions to a client.
+
+    Args:
+        site_id: Site identifier
+        client_mac: Client MAC address
+        settings: Application settings
+        upload_limit_kbps: Upload speed limit in kbps
+        download_limit_kbps: Download speed limit in kbps
+        confirm: Confirmation flag (must be True to execute)
+        dry_run: If True, validate but don't apply limits
+
+    Returns:
+        Bandwidth limit result dictionary
+
+    Raises:
+        ConfirmationRequiredError: If confirm is not True
+    """
+    site_id = validate_site_id(site_id)
+    client_mac = validate_mac_address(client_mac)
+    validate_confirmation(confirm, "client management operation")
+    logger = get_logger(__name__, settings.log_level)
+
+    # Validate bandwidth limits
+    if upload_limit_kbps is not None and upload_limit_kbps <= 0:
+        raise ValueError("Upload limit must be positive")
+    if download_limit_kbps is not None and download_limit_kbps <= 0:
+        raise ValueError("Download limit must be positive")
+
+    parameters = {
+        "site_id": site_id,
+        "client_mac": client_mac,
+        "upload_limit_kbps": upload_limit_kbps,
+        "download_limit_kbps": download_limit_kbps,
+    }
+
+    if dry_run:
+        logger.info(
+            sanitize_log_message(
+                f"DRY RUN: Would apply bandwidth limits to client '{client_mac}' in site '{site_id}'"
+            )
+        )
+        log_audit(
+            operation="limit_bandwidth",
+            parameters=parameters,
+            result="dry_run",
+            site_id=site_id,
+            dry_run=True,
+        )
+        return {
+            "dry_run": True,
+            "would_limit": client_mac,
+            "upload_limit_kbps": upload_limit_kbps,
+            "download_limit_kbps": download_limit_kbps,
+        }
+
+    try:
+        async with UniFiClient(settings) as client:
+            await client.authenticate()
+
+            # Build bandwidth limit payload
+            limit_data = {
+                "action": "limit-bandwidth",
+                "params": {},
+            }
+
+            if upload_limit_kbps is not None:
+                limit_data["params"]["uploadLimit"] = upload_limit_kbps  # type: ignore[index]
+            if download_limit_kbps is not None:
+                limit_data["params"]["downloadLimit"] = download_limit_kbps  # type: ignore[index]
+
+            # Apply bandwidth limits using new API endpoint
+            await client.post(
+                f"/integration/v1/sites/{site_id}/clients/{client_mac}/action",
+                json_data=limit_data,
+            )
+
+            logger.info(
+                sanitize_log_message(
+                    f"Applied bandwidth limits to client '{client_mac}' in site '{site_id}'"
+                )
+            )
+            log_audit(
+                operation="limit_bandwidth",
+                parameters=parameters,
+                result="success",
+                site_id=site_id,
+            )
+
+            return {
+                "success": True,
+                "client_mac": client_mac,
+                "upload_limit_kbps": upload_limit_kbps,
+                "download_limit_kbps": download_limit_kbps,
+                "message": "Bandwidth limits applied",
+            }
+
+    except Exception as e:
+        logger.error(
+            sanitize_log_message(f"Failed to apply bandwidth limits to client '{client_mac}': {e}")
+        )
+        log_audit(
+            operation="limit_bandwidth",
+            parameters=parameters,
+            result="failed",
+            site_id=site_id,
+        )
+        raise

src/tools/clients.py

@@ -0,0 +1,203 @@
+"""Client management MCP tools."""
+
+from typing import Any
+
+from ..api import UniFiClient
+from ..config import Settings
+from ..models import Client
+from ..utils import (
+    ResourceNotFoundError,
+    get_logger,
+    sanitize_log_message,
+    validate_limit_offset,
+    validate_mac_address,
+    validate_site_id,
+)
+
+
+async def get_client_details(site_id: str, client_mac: str, settings: Settings) -> dict[str, Any]:
+    """Get detailed information for a specific client.
+
+    Args:
+        site_id: Site identifier
+        client_mac: Client MAC address
+        settings: Application settings
+
+    Returns:
+        Client details dictionary
+
+    Raises:
+        ResourceNotFoundError: If client not found
+    """
+    site_id = validate_site_id(site_id)
+    client_mac = validate_mac_address(client_mac)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        # Try active clients first
+        response = await client.get(f"/ea/sites/{site_id}/sta")
+        clients_data = response.get("data", []) if isinstance(response, dict) else response
+
+        for client_data in clients_data:
+            if validate_mac_address(client_data.get("mac", "")) == client_mac:
+                client_obj = Client(**client_data)
+                logger.info(sanitize_log_message(f"Retrieved client details for {client_mac}"))
+                return client_obj.model_dump()  # type: ignore[no-any-return]
+
+        # If not found in active, try all users
+        response = await client.get(f"/ea/sites/{site_id}/stat/alluser")
+        clients_data = response.get("data", []) if isinstance(response, dict) else response
+
+        for client_data in clients_data:
+            if validate_mac_address(client_data.get("mac", "")) == client_mac:
+                client_obj = Client(**client_data)
+                logger.info(sanitize_log_message(f"Retrieved client details for {client_mac}"))
+                return client_obj.model_dump()  # type: ignore[no-any-return]
+
+        raise ResourceNotFoundError("client", client_mac)
+
+
+async def get_client_statistics(
+    site_id: str, client_mac: str, settings: Settings
+) -> dict[str, Any]:
+    """Retrieve bandwidth and connection statistics for a client.
+
+    Args:
+        site_id: Site identifier
+        client_mac: Client MAC address
+        settings: Application settings
+
+    Returns:
+        Client statistics dictionary
+
+    Raises:
+        ResourceNotFoundError: If client not found
+    """
+    site_id = validate_site_id(site_id)
+    client_mac = validate_mac_address(client_mac)
+    logger = get_logger(__name__, settings.log_level)
+
+    async with UniFiClient(settings) as client:
+        await client.authenticate()
+
+        # Get from active clients
+        response = await client.get(f"/ea/sites/{site_id}/sta")
+        clients_data = response.get("data", []) if isinstance(response, dict) else response
+
+        for client_data in clients_data:
+            if validate_mac_address(client_data.get("mac", "")) == client_mac:
+                # Extract statistics
+                stats = {
+                    "mac": client_mac,
+                    "tx_bytes": client_data.get("tx_bytes", 0),
+                    "rx_bytes": client_data.get("rx_bytes", 0),
+                    "tx_packets": client_data.get("tx_packets", 0),
+                    "rx_packets": client_data.get("rx_packets", 0),
+                    "tx_rate": client_data.get("tx_rate"),
+                    "rx_rate": client_data.get("rx_rate"),
+                    "signal": client_data.get("signal"),
+                    "rssi": client_data.get("rssi"),
+                    "noise": client_data.get("noise"),
+                    "uptime": client_data.get("uptime", 0),
+                    "is_wired": client_data.get("is_wired", False),
+                }
+                logger.info(sanitize_log_message(f"Retrieved statistics for client {client_mac}"))
+                return stats
+
+        raise ResourceNotFoundError("client", client_mac)
+
+
+async def list_active_clients(
+    site_id: str,
+    settings: Settings,
+    limit: int | None = None,
+    offset: int | None = None,
+) -> list[dict[str, Any]]:
+    """List currently connected clients.
+
+    Args:
+        site_id: Site identifier
+        settings: Application settings
+        limit: Maximum number of clients to return
+        offset: Number of clients to skip
+
+    Returns:
+        List of active client 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}/sta")
+        clients_data = response.get("data", []) if isinstance(response, dict) else response
+
+        # Apply pagination
+        paginated = clients_data[offset : offset + limit]
+
+        # Parse into Client models
+        clients = [Client(**c).model_dump() for c in paginated]
+
+        logger.info(
+            sanitize_log_message(f"Retrieved {len(clients)} active clients for site '{site_id}'")
+        )
+        return clients
+
+
+async def search_clients(
+    site_id: str,
+    query: str,
+    settings: Settings,
+    limit: int | None = None,
+    offset: int | None = None,
+) -> list[dict[str, Any]]:
+    """Search clients by MAC, IP, or hostname.
+
+    Args:
+        site_id: Site identifier
+        query: Search query string
+        settings: Application settings
+        limit: Maximum number of clients to return
+        offset: Number of clients to skip
+
+    Returns:
+        List of matching client 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()
+
+        # Search in all users
+        response = await client.get(f"/ea/sites/{site_id}/stat/alluser")
+        clients_data = response.get("data", []) if isinstance(response, dict) else response
+
+        # Search by MAC, IP, hostname, or name
+        query_lower = query.lower()
+        filtered = [
+            c
+            for c in clients_data
+            if query_lower in c.get("mac", "").lower()
+            or query_lower in c.get("ip", "").lower()
+            or query_lower in c.get("hostname", "").lower()
+            or query_lower in c.get("name", "").lower()
+        ]
+
+        # Apply pagination
+        paginated = filtered[offset : offset + limit]
+
+        # Parse into Client models
+        clients = [Client(**c).model_dump() for c in paginated]
+
+        logger.info(
+            sanitize_log_message(
+                f"Found {len(clients)} clients matching '{query}' in site '{site_id}'"
+            )
+        )
+        return clients