Megaraptor MCP

Instructions

Implementation Reference

• src/megaraptor_mcp/tools/clients.py:172-237 (handler)

  The tool 'label_client' is registered with '@mcp.tool()' and implements label management using Velociraptor VQL.

  @mcp.tool()
  async def label_client(
      client_id: str,
      labels: list[str],
      operation: str = "add",
  ) -> list[TextContent]:
      """Add or remove labels from a Velociraptor client.
  
      Args:
          client_id: The client ID (e.g., 'C.1234567890abcdef')
          labels: List of label names to add or remove
          operation: Either 'add' or 'remove' (default: 'add')
  
      Returns:
          Updated client labels.
      """
      try:
          # Validate client_id
          client_id = validate_client_id(client_id)
  
          if operation not in ("add", "remove"):
              return [TextContent(
                  type="text",
                  text=json.dumps({
                      "error": "Operation must be 'add' or 'remove'",
                      "hint": "Use operation='add' to add labels or operation='remove' to remove them"
                  })
              )]
  
          client = get_client()
  
          # Build the VQL for label modification
          labels_str = ", ".join(f"'{label}'" for label in labels)
  
          if operation == "add":
              vql = f"SELECT label(client_id='{client_id}', labels=[{labels_str}], op='set') FROM scope()"
          else:
              vql = f"SELECT label(client_id='{client_id}', labels=[{labels_str}], op='remove') FROM scope()"
  
          results = client.query(vql)
  
          # Get updated client info
          info_vql = f"SELECT labels FROM clients(client_id='{client_id}')"
          info_results = client.query(info_vql)
  
          return [TextContent(
              type="text",
              text=json.dumps({
                  "client_id": client_id,
                  "operation": operation,
                  "labels_modified": labels,
                  "current_labels": info_results[0].get("labels", []) if info_results else [],
              }, indent=2)
          )]
  
      except ValueError as e:
          # Validation errors
          return [TextContent(
              type="text",
              text=json.dumps({
                  "error": str(e),
                  "hint": "Provide a valid client ID starting with 'C.'"
              })
          )]
  
      except grpc.RpcError as e: