go2tool 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,5 @@
1
+ __pycache__/
2
+ *.pyc
3
+ .venv/
4
+ dist/
5
+ uv.lock
go2tool-0.1.0/LICENSE ADDED
@@ -0,0 +1,17 @@
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ Copyright 2025 Dimensional Inc.
6
+
7
+ Licensed under the Apache License, Version 2.0 (the "License");
8
+ you may not use this file except in compliance with the License.
9
+ You may obtain a copy of the License at
10
+
11
+ http://www.apache.org/licenses/LICENSE-2.0
12
+
13
+ Unless required by applicable law or agreed to in writing, software
14
+ distributed under the License is distributed on an "AS IS" BASIS,
15
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
+ See the License for the specific language governing permissions and
17
+ limitations under the License.
go2tool-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,79 @@
1
+ Metadata-Version: 2.5
2
+ Name: go2tool
3
+ Version: 0.1.0
4
+ Summary: Set up Unitree Go2 robots: Bluetooth Wi-Fi provisioning and LAN/BLE discovery
5
+ Project-URL: Homepage, https://github.com/jeff-hykin/go2tool
6
+ Project-URL: Repository, https://github.com/jeff-hykin/go2tool
7
+ Author: Dimensional Inc.
8
+ License-Expression: Apache-2.0
9
+ License-File: LICENSE
10
+ Keywords: bluetooth,go2,robot,unitree,wifi
11
+ Classifier: Operating System :: MacOS
12
+ Classifier: Operating System :: POSIX :: Linux
13
+ Classifier: Programming Language :: Python :: 3
14
+ Requires-Python: >=3.10
15
+ Requires-Dist: bleak>=1.0.0
16
+ Requires-Dist: cryptography>=42
17
+ Requires-Dist: psutil>=7.0.0
18
+ Requires-Dist: typer<1,>=0.19.2
19
+ Description-Content-Type: text/markdown
20
+
21
+ # go2tool
22
+
23
+ Set up a Unitree Go2 (Pro or Air, stock firmware) from your laptop: put it on Wi-Fi over Bluetooth and find its IP. No Unitree app, no jailbreak.
24
+
25
+ Originally the `dimos go2tool` command from [dimOS](https://github.com/dimensionalOS/dimos); this is the same tool as a standalone package.
26
+
27
+ ## Install
28
+
29
+ ```bash
30
+ pip install go2tool
31
+ # or run it without installing
32
+ uvx go2tool discover
33
+ ```
34
+
35
+ Needs Bluetooth on the machine running it (macOS or Linux/BlueZ) for the BLE parts.
36
+
37
+ ## Usage
38
+
39
+ 1. Power on the Go2. It advertises over BLE immediately.
40
+
41
+ 2. (Optional) make sure the robot is detected:
42
+
43
+ ```bash
44
+ go2tool discover
45
+ ```
46
+
47
+ 3. Provision Wi-Fi (one-time per network):
48
+
49
+ ```bash
50
+ go2tool connect-wifi --ssid <wifi> --password <password>
51
+ ```
52
+
53
+ Scans BLE and connects to the only robot it finds, or prompts you to pick if there are several. Use `--name Go2_49060`, `--serial <serial>` or `--mac <ble-mac>` to pick one non-interactively, and `--ap` to make the robot host a hotspot with that SSID/password instead of joining one.
54
+
55
+ 4. Find the robot's IP:
56
+
57
+ ```bash
58
+ go2tool discover
59
+ ```
60
+
61
+ Prints `SOURCE NAME IP MAC SERIAL` for every robot it sees over BLE and LAN, e.g.
62
+
63
+ ```
64
+ SOURCE NAME IP MAC SERIAL
65
+ BLE Go2_49060 - AA:BB:CC:DD:EE:FF B42D4000P6PC49060
66
+ LAN - 192.168.1.42 - B42D4000P6PC49060
67
+ ```
68
+
69
+ Options: `--lan` / `--ble` to use only one source, `--timeout 0` to keep listening, `--lan-tick` for the LAN poll interval.
70
+
71
+ LAN discovery pins the UDP multicast probe to each real network interface, so it still works with Tailscale or other VPNs up.
72
+
73
+ ## Credits
74
+
75
+ BLE provisioning protocol from the [UniPwn](https://github.com/Bin4ry/UniPwn) project's reverse engineering (exploit payload removed; only the normal provisioning flow the official app uses).
76
+
77
+ ## License
78
+
79
+ Apache-2.0, Copyright 2025-2026 Dimensional Inc.
@@ -0,0 +1,59 @@
1
+ # go2tool
2
+
3
+ Set up a Unitree Go2 (Pro or Air, stock firmware) from your laptop: put it on Wi-Fi over Bluetooth and find its IP. No Unitree app, no jailbreak.
4
+
5
+ Originally the `dimos go2tool` command from [dimOS](https://github.com/dimensionalOS/dimos); this is the same tool as a standalone package.
6
+
7
+ ## Install
8
+
9
+ ```bash
10
+ pip install go2tool
11
+ # or run it without installing
12
+ uvx go2tool discover
13
+ ```
14
+
15
+ Needs Bluetooth on the machine running it (macOS or Linux/BlueZ) for the BLE parts.
16
+
17
+ ## Usage
18
+
19
+ 1. Power on the Go2. It advertises over BLE immediately.
20
+
21
+ 2. (Optional) make sure the robot is detected:
22
+
23
+ ```bash
24
+ go2tool discover
25
+ ```
26
+
27
+ 3. Provision Wi-Fi (one-time per network):
28
+
29
+ ```bash
30
+ go2tool connect-wifi --ssid <wifi> --password <password>
31
+ ```
32
+
33
+ Scans BLE and connects to the only robot it finds, or prompts you to pick if there are several. Use `--name Go2_49060`, `--serial <serial>` or `--mac <ble-mac>` to pick one non-interactively, and `--ap` to make the robot host a hotspot with that SSID/password instead of joining one.
34
+
35
+ 4. Find the robot's IP:
36
+
37
+ ```bash
38
+ go2tool discover
39
+ ```
40
+
41
+ Prints `SOURCE NAME IP MAC SERIAL` for every robot it sees over BLE and LAN, e.g.
42
+
43
+ ```
44
+ SOURCE NAME IP MAC SERIAL
45
+ BLE Go2_49060 - AA:BB:CC:DD:EE:FF B42D4000P6PC49060
46
+ LAN - 192.168.1.42 - B42D4000P6PC49060
47
+ ```
48
+
49
+ Options: `--lan` / `--ble` to use only one source, `--timeout 0` to keep listening, `--lan-tick` for the LAN poll interval.
50
+
51
+ LAN discovery pins the UDP multicast probe to each real network interface, so it still works with Tailscale or other VPNs up.
52
+
53
+ ## Credits
54
+
55
+ BLE provisioning protocol from the [UniPwn](https://github.com/Bin4ry/UniPwn) project's reverse engineering (exploit payload removed; only the normal provisioning flow the official app uses).
56
+
57
+ ## License
58
+
59
+ Apache-2.0, Copyright 2025-2026 Dimensional Inc.
@@ -0,0 +1,32 @@
1
+ [project]
2
+ name = "go2tool"
3
+ version = "0.1.0"
4
+ description = "Set up Unitree Go2 robots: Bluetooth Wi-Fi provisioning and LAN/BLE discovery"
5
+ readme = "README.md"
6
+ license = "Apache-2.0"
7
+ license-files = ["LICENSE"]
8
+ authors = [{ name = "Dimensional Inc." }]
9
+ requires-python = ">=3.10"
10
+ dependencies = [
11
+ "typer>=0.19.2,<1",
12
+ "psutil>=7.0.0",
13
+ "bleak>=1.0.0",
14
+ "cryptography>=42",
15
+ ]
16
+ keywords = ["unitree", "go2", "robot", "bluetooth", "wifi"]
17
+ classifiers = [
18
+ "Programming Language :: Python :: 3",
19
+ "Operating System :: MacOS",
20
+ "Operating System :: POSIX :: Linux",
21
+ ]
22
+
23
+ [project.urls]
24
+ Homepage = "https://github.com/jeff-hykin/go2tool"
25
+ Repository = "https://github.com/jeff-hykin/go2tool"
26
+
27
+ [project.scripts]
28
+ go2tool = "go2tool.cli:main"
29
+
30
+ [build-system]
31
+ requires = ["hatchling>=1.27"]
32
+ build-backend = "hatchling.build"
@@ -0,0 +1,17 @@
1
+ # Copyright 2025-2026 Dimensional Inc.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """Unitree Go2 setup utilities: BLE Wi-Fi provisioning and network discovery."""
16
+
17
+ __version__ = "0.1.0"
@@ -0,0 +1,355 @@
1
+ # Copyright 2025-2026 Dimensional Inc.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """BLE wifi-provisioning protocol for Unitree robots.
16
+
17
+ Protocol reverse-engineered by the UniPwn project (https://github.com/Bin4ry/UniPwn);
18
+ this module strips the exploit payload and keeps only the legitimate provisioning
19
+ flow used by the official Unitree app.
20
+ """
21
+
22
+ from __future__ import annotations
23
+
24
+ import asyncio
25
+ from collections.abc import AsyncIterator, Awaitable, Callable
26
+ from dataclasses import dataclass
27
+ from typing import Any
28
+
29
+ UNITREE_NAME_PREFIXES = ("Go2_", "G1_", "B2_", "H1_", "X1_")
30
+
31
+ UNITREE_SERVICE_UUID = "0000ffe0-0000-1000-8000-00805f9b34fb"
32
+ NOTIFY_CHAR_UUID = "0000ffe1-0000-1000-8000-00805f9b34fb"
33
+ WRITE_CHAR_UUID = "0000ffe2-0000-1000-8000-00805f9b34fb"
34
+
35
+ # Symmetric AES-CFB-128 key/iv burned into the firmware.
36
+ _AES_KEY = bytes.fromhex("df98b715d5c6ed2b25817b6f2554124a")
37
+ _AES_IV = bytes.fromhex("2841ae97419c2973296a0d4bdfe19a4f")
38
+
39
+ CHUNK_SIZE = 14
40
+ HANDSHAKE_CONTENT = b"unitree"
41
+
42
+ # Instruction opcodes (TX / RX use the same numbers; opcode byte differs: 0x52 / 0x51).
43
+ INST_HANDSHAKE = 1
44
+ INST_SERIAL = 2
45
+ INST_INIT_STA = 3
46
+ INST_SSID = 4
47
+ INST_PASSWORD = 5
48
+ INST_COUNTRY = 6
49
+
50
+
51
+ @dataclass
52
+ class Go2Device:
53
+ name: str
54
+ address: str
55
+ serial: str | None = None
56
+
57
+
58
+ def _serial_from_manuf(manuf_data: dict[int, bytes] | None) -> str | None:
59
+ """Recover the full 16-char serial from BLE manufacturer data.
60
+
61
+ The Go2 packs its serial across the company-ID + payload field: the
62
+ company-ID's little-endian bytes spell the first two ASCII chars, the
63
+ payload spells the rest.
64
+ """
65
+ if not manuf_data:
66
+ return None
67
+ cid, payload = next(iter(manuf_data.items()))
68
+ prefix = bytes([cid & 0xFF, (cid >> 8) & 0xFF])
69
+ try:
70
+ return (prefix + payload).decode("ascii")
71
+ except UnicodeDecodeError:
72
+ return None
73
+
74
+
75
+ def _cipher() -> Any:
76
+ from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
77
+
78
+ return Cipher(algorithms.AES(_AES_KEY), modes.CFB(_AES_IV))
79
+
80
+
81
+ def encrypt(data: bytes) -> bytes:
82
+ enc = _cipher().encryptor()
83
+ out: bytes = enc.update(data) + enc.finalize()
84
+ return out
85
+
86
+
87
+ def decrypt(data: bytes) -> bytes:
88
+ dec = _cipher().decryptor()
89
+ out: bytes = dec.update(data) + dec.finalize()
90
+ return out
91
+
92
+
93
+ def build_packet(instruction: int, payload: bytes = b"") -> bytes:
94
+ """Build an encrypted TX packet: 0x52, length, instruction, *payload, checksum."""
95
+ body = bytes([0x52, len(payload) + 4, instruction]) + payload
96
+ checksum = (-sum(body)) & 0xFF
97
+ return encrypt(body + bytes([checksum]))
98
+
99
+
100
+ def validate_response(response: bytes, expected_inst: int) -> bool:
101
+ if len(response) < 5:
102
+ return False
103
+ if response[0] != 0x51:
104
+ return False
105
+ if len(response) != response[1]:
106
+ return False
107
+ if response[2] != expected_inst:
108
+ return False
109
+ if (-sum(response[:-1])) & 0xFF != response[-1]:
110
+ return False
111
+ return response[3] == 0x01
112
+
113
+
114
+ async def _stop_stale_bluez_scan() -> None:
115
+ """Tell BlueZ to drop any leftover scan from a prior process.
116
+
117
+ `BleakScanner.__aexit__` is what normally calls StopDiscovery. If a previous
118
+ invocation was killed before that ran (e.g. SIGTERM from `timeout`), BlueZ
119
+ keeps the session and the next StartDiscovery returns `InProgress`. Best-effort.
120
+ """
121
+ try:
122
+ proc = await asyncio.create_subprocess_exec(
123
+ "bluetoothctl",
124
+ "scan",
125
+ "off",
126
+ stdout=asyncio.subprocess.DEVNULL,
127
+ stderr=asyncio.subprocess.DEVNULL,
128
+ )
129
+ try:
130
+ await asyncio.wait_for(proc.wait(), timeout=2.0)
131
+ except asyncio.TimeoutError:
132
+ proc.kill()
133
+ except (FileNotFoundError, OSError):
134
+ pass
135
+
136
+
137
+ async def discover_ble(
138
+ prefixes: tuple[str, ...] = UNITREE_NAME_PREFIXES,
139
+ ) -> AsyncIterator[Go2Device]:
140
+ """Stream Unitree robots seen over BLE. Runs until the consumer stops iterating.
141
+
142
+ Yields a device on first sighting, and again whenever a later adv packet
143
+ upgrades a previously-unknown serial number.
144
+ """
145
+ from bleak import BleakScanner
146
+
147
+ await _stop_stale_bluez_scan()
148
+
149
+ found: dict[str, Go2Device] = {}
150
+ queue: asyncio.Queue[Go2Device] = asyncio.Queue()
151
+
152
+ def _on_detect(device: Any, adv: Any) -> None:
153
+ if not device.name or not device.name.startswith(prefixes):
154
+ return
155
+ serial = _serial_from_manuf(getattr(adv, "manufacturer_data", None))
156
+ existing = found.get(device.address)
157
+ if existing is None:
158
+ dev = Go2Device(name=device.name, address=device.address, serial=serial)
159
+ found[device.address] = dev
160
+ queue.put_nowait(dev)
161
+ elif existing.serial is None and serial is not None:
162
+ existing.serial = serial
163
+ queue.put_nowait(existing)
164
+
165
+ async with BleakScanner(detection_callback=_on_detect):
166
+ while True:
167
+ yield await queue.get()
168
+
169
+
170
+ async def find_robots(
171
+ timeout: float = 15.0,
172
+ prefixes: tuple[str, ...] = UNITREE_NAME_PREFIXES,
173
+ on_device: Callable[[Go2Device], None] | None = None,
174
+ ) -> list[Go2Device]:
175
+ """One-shot BLE scan: collects devices for `timeout` seconds and returns them."""
176
+ out: dict[str, Go2Device] = {}
177
+
178
+ async def _collect() -> None:
179
+ async for dev in discover_ble(prefixes=prefixes):
180
+ out[dev.address] = dev
181
+ if on_device is not None:
182
+ on_device(dev)
183
+
184
+ try:
185
+ await asyncio.wait_for(_collect(), timeout=timeout)
186
+ except asyncio.TimeoutError:
187
+ pass
188
+ return list(out.values())
189
+
190
+
191
+ class _Session:
192
+ """Internal helper that wraps a BleakClient with the unitree request/response state."""
193
+
194
+ def __init__(self, client: Any) -> None:
195
+ self.client = client
196
+ self.event = asyncio.Event()
197
+ self.last: bytes | None = None
198
+ self.serial_chunks: dict[int, bytes] = {}
199
+ self.serial: bytes | None = None
200
+
201
+ def on_notify(self, _sender: Any, data: bytearray) -> None:
202
+ packet = decrypt(bytes(data))
203
+ if len(packet) < 5 or packet[0] != 0x51:
204
+ return
205
+ if packet[2] == INST_SERIAL:
206
+ idx, total = packet[3], packet[4]
207
+ self.serial_chunks[idx] = packet[5:-1]
208
+ if len(self.serial_chunks) >= total:
209
+ self.serial = b"".join(self.serial_chunks[i] for i in sorted(self.serial_chunks))
210
+ self.event.set()
211
+ else:
212
+ self.last = packet
213
+ self.event.set()
214
+
215
+ async def write(self, packet: bytes) -> None:
216
+ await self.client.write_gatt_char(WRITE_CHAR_UUID, packet, response=True)
217
+
218
+ async def write_validated(
219
+ self, packet: bytes, expected_inst: int, timeout: float = 10.0
220
+ ) -> None:
221
+ await self.write(packet)
222
+ await asyncio.wait_for(self.event.wait(), timeout)
223
+ if self.last is None or not validate_response(self.last, expected_inst):
224
+ raise RuntimeError(f"BLE response invalid for instruction {expected_inst}")
225
+ self.event.clear()
226
+ self.last = None
227
+
228
+ async def send_chunked(
229
+ self, instruction: int, data: bytes, response_timeout: float = 10.0
230
+ ) -> None:
231
+ total = max(1, (len(data) + CHUNK_SIZE - 1) // CHUNK_SIZE)
232
+ for i in range(total):
233
+ chunk = data[i * CHUNK_SIZE : (i + 1) * CHUNK_SIZE]
234
+ packet = build_packet(instruction, bytes([i + 1, total]) + chunk)
235
+ await self.write(packet)
236
+ await asyncio.sleep(0.1)
237
+ await asyncio.wait_for(self.event.wait(), response_timeout)
238
+ if self.last is None or not validate_response(self.last, instruction):
239
+ raise RuntimeError(f"BLE response invalid for chunked instruction {instruction}")
240
+ self.event.clear()
241
+ self.last = None
242
+
243
+
244
+ async def _connect_with_retry(
245
+ address: str,
246
+ timeout: float,
247
+ attempts: int,
248
+ on_progress: Callable[[str], None],
249
+ ) -> Any:
250
+ """Open a BleakClient, retrying only the connection step."""
251
+ from bleak import BleakClient
252
+
253
+ last_exc: BaseException | None = None
254
+ for i in range(attempts):
255
+ client = BleakClient(address, timeout=timeout)
256
+ try:
257
+ await client.connect()
258
+ return client
259
+ except Exception as e:
260
+ last_exc = e
261
+ on_progress(f"connect attempt {i + 1}/{attempts} failed: {e}")
262
+ if i + 1 < attempts:
263
+ await asyncio.sleep(1.0)
264
+ assert last_exc is not None
265
+ raise last_exc
266
+
267
+
268
+ async def provision_wifi(
269
+ address: str,
270
+ ssid: str,
271
+ password: str,
272
+ country_code: str = "US",
273
+ *,
274
+ ap: bool = False,
275
+ timeout: float = 30.0,
276
+ connect_retries: int = 3,
277
+ on_progress: Callable[[str], None] | None = None,
278
+ ) -> str | None:
279
+ """Provision a Unitree robot's wifi over BLE. Returns the serial number on success.
280
+
281
+ With `ap=True` the robot hosts its own hotspot named `ssid` instead of joining one.
282
+
283
+ Retries the BLE connection step up to `connect_retries` times. Once connected,
284
+ handshake/SSID/password/country failures are not retried — those indicate a
285
+ protocol-level problem (or partial state on the robot) where blind retry is
286
+ counterproductive.
287
+ """
288
+ progress = on_progress or (lambda _msg: None)
289
+
290
+ client = await _connect_with_retry(address, timeout, connect_retries, progress)
291
+ try:
292
+ session = _Session(client)
293
+ try:
294
+ await client.start_notify(NOTIFY_CHAR_UUID, session.on_notify)
295
+ except Exception:
296
+ # Some firmware versions only expose the raw GATT handle.
297
+ await client.start_notify(13, session.on_notify) # type: ignore[arg-type]
298
+
299
+ progress("handshake")
300
+ await session.write_validated(
301
+ build_packet(INST_HANDSHAKE, bytes([0, 0]) + HANDSHAKE_CONTENT),
302
+ INST_HANDSHAKE,
303
+ )
304
+
305
+ progress("read serial")
306
+ await session.write(build_packet(INST_SERIAL, bytes([0])))
307
+ try:
308
+ await asyncio.wait_for(session.event.wait(), 2.0)
309
+ except asyncio.TimeoutError:
310
+ pass
311
+ session.event.clear()
312
+
313
+ progress(f"init {'AP' if ap else 'STA'} mode")
314
+ await session.write_validated(
315
+ build_packet(INST_INIT_STA, bytes([1 if ap else 2])), INST_INIT_STA
316
+ )
317
+
318
+ progress(f"set SSID: {ssid}")
319
+ await session.send_chunked(INST_SSID, ssid.encode("utf-8"))
320
+
321
+ progress("set password")
322
+ await session.send_chunked(INST_PASSWORD, password.encode("utf-8"), response_timeout=5.0)
323
+
324
+ progress(f"set country: {country_code}")
325
+ await session.write_validated(
326
+ build_packet(INST_COUNTRY, bytes([1]) + country_code.encode("utf-8") + b"\x00"),
327
+ INST_COUNTRY,
328
+ )
329
+ finally:
330
+ await client.disconnect()
331
+
332
+ if session.serial:
333
+ return session.serial.decode("utf-8", errors="replace").rstrip("\x00")
334
+ return None
335
+
336
+
337
+ async def retry(
338
+ fn: Callable[[], Awaitable[Any]],
339
+ attempts: int = 3,
340
+ delay: float = 1.0,
341
+ on_error: Callable[[int, BaseException], None] | None = None,
342
+ ) -> Any:
343
+ """Retry an awaitable a few times. Re-raises the final exception."""
344
+ last_exc: BaseException | None = None
345
+ for i in range(attempts):
346
+ try:
347
+ return await fn()
348
+ except Exception as e:
349
+ last_exc = e
350
+ if on_error:
351
+ on_error(i + 1, e)
352
+ if i + 1 < attempts:
353
+ await asyncio.sleep(delay)
354
+ assert last_exc is not None
355
+ raise last_exc
@@ -0,0 +1,209 @@
1
+ # Copyright 2025-2026 Dimensional Inc.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """`go2tool` — Go2 setup utilities (BLE wifi provisioning, network discovery)."""
16
+
17
+ from __future__ import annotations
18
+
19
+ import asyncio
20
+
21
+ import typer
22
+
23
+ app = typer.Typer(
24
+ help="Go2 setup utilities (BLE wifi provisioning, network discovery)",
25
+ no_args_is_help=True,
26
+ )
27
+
28
+
29
+ _HEADER = f"{'SOURCE':<6} {'NAME':<14} {'IP':<15} {'MAC':<19} SERIAL"
30
+
31
+
32
+ def _format_row(source: str, name: str, ip: str, mac: str, serial: str) -> str:
33
+ return f"{source:<6} {name:<14} {ip:<15} {mac:<19} {serial}"
34
+
35
+
36
+ @app.command("discover")
37
+ def discover(
38
+ ble: bool = typer.Option(False, "--ble", help="BLE only (default: BLE + LAN)"),
39
+ lan: bool = typer.Option(False, "--lan", help="LAN only (default: BLE + LAN)"),
40
+ lan_tick: float = typer.Option(2.0, "--lan-tick", help="LAN poll interval (s)"),
41
+ timeout: float = typer.Option(
42
+ 7.0, "--timeout", "-t", help="Stop after this many seconds (0 = run forever)"
43
+ ),
44
+ ) -> None:
45
+ """Stream Go2 robot discoveries from BLE and/or LAN."""
46
+ do_ble = ble or not lan
47
+ do_lan = lan or not ble
48
+
49
+ typer.echo(_HEADER)
50
+
51
+ import signal
52
+
53
+ async def run() -> None:
54
+ from go2tool.ble import discover_ble
55
+ from go2tool.landiscovery import discover_lan
56
+
57
+ seen_ble: set[tuple[str, str | None]] = set()
58
+ seen_lan: set[str] = set()
59
+
60
+ async def _consume_ble() -> None:
61
+ async for d in discover_ble():
62
+ key = (d.address, d.serial)
63
+ if key in seen_ble:
64
+ continue
65
+ seen_ble.add(key)
66
+ typer.echo(_format_row("BLE", d.name, "-", d.address, d.serial or "?"))
67
+
68
+ async def _consume_lan() -> None:
69
+ async for d in discover_lan(tick=lan_tick):
70
+ if d.serial in seen_lan:
71
+ continue
72
+ seen_lan.add(d.serial)
73
+ typer.echo(_format_row("LAN", "-", d.ip, d.mac or "-", d.serial))
74
+
75
+ tasks: list[asyncio.Task[None]] = []
76
+ if do_ble:
77
+ tasks.append(asyncio.create_task(_consume_ble()))
78
+ if do_lan:
79
+ tasks.append(asyncio.create_task(_consume_lan()))
80
+
81
+ # Cancel on SIGINT or SIGTERM so the BleakScanner's __aexit__ runs and
82
+ # calls StopDiscovery; otherwise BlueZ retains the scan session.
83
+ loop = asyncio.get_running_loop()
84
+
85
+ def _stop() -> None:
86
+ for t in tasks:
87
+ t.cancel()
88
+
89
+ for sig in (signal.SIGINT, signal.SIGTERM):
90
+ loop.add_signal_handler(sig, _stop)
91
+
92
+ if timeout > 0:
93
+ loop.call_later(timeout, _stop)
94
+
95
+ await asyncio.gather(*tasks, return_exceptions=True)
96
+
97
+ asyncio.run(run())
98
+ typer.echo("\nStopped.")
99
+
100
+
101
+ @app.command("connect-wifi")
102
+ def connect_wifi(
103
+ ssid: str | None = typer.Option(None, "--ssid", help="Wi-Fi SSID"),
104
+ password: str | None = typer.Option(None, "--password", help="Wi-Fi password"),
105
+ country: str = typer.Option("US", "--country", help="Two-letter country code"),
106
+ ap: bool = typer.Option(False, "--ap", help="Host a hotspot with --ssid/--password instead"),
107
+ mac: str | None = typer.Option(None, "--mac", help="BLE MAC (skip scan)"),
108
+ serial: str | None = typer.Option(
109
+ None, "--serial", help="Robot serial — scan and auto-select match"
110
+ ),
111
+ name: str | None = typer.Option(
112
+ None, "--name", help="Robot BLE name (e.g. Go2_49060) — scan and auto-select match"
113
+ ),
114
+ timeout: float = typer.Option(5.0, "--timeout", help="Scan / connect timeout in seconds"),
115
+ retries: int = typer.Option(3, "--retries", help="Number of provisioning attempts"),
116
+ ) -> None:
117
+ """Provision a Go2 with Wi-Fi credentials over Bluetooth.
118
+
119
+ Fully non-interactive when (--mac | --serial | --name) and --ssid/--password
120
+ are all provided.
121
+ """
122
+ from go2tool.ble import (
123
+ Go2Device,
124
+ find_robots,
125
+ provision_wifi,
126
+ retry,
127
+ )
128
+
129
+ async def run() -> None:
130
+ if mac is not None:
131
+ target = mac
132
+ else:
133
+ typer.echo(f"Scanning BLE for {timeout:.0f}s ...")
134
+
135
+ def _on_device(d: Go2Device) -> None:
136
+ typer.echo(_format_row("BLE", d.name, "-", d.address, d.serial or "?"))
137
+
138
+ devices = await find_robots(timeout=timeout, on_device=_on_device)
139
+ if not devices:
140
+ typer.echo("No Unitree robots detected.", err=True)
141
+ raise typer.Exit(1)
142
+
143
+ if serial is not None or name is not None:
144
+ matches = [
145
+ d
146
+ for d in devices
147
+ if (serial is None or d.serial == serial) and (name is None or d.name == name)
148
+ ]
149
+ if not matches:
150
+ crit = ", ".join(
151
+ f"{k}={v}" for k, v in (("serial", serial), ("name", name)) if v is not None
152
+ )
153
+ typer.echo(f"No BLE device matched {crit}.", err=True)
154
+ raise typer.Exit(1)
155
+ if len(matches) > 1:
156
+ typer.echo(f"{len(matches)} devices match — refine criteria:", err=True)
157
+ for d in matches:
158
+ typer.echo(f" {d.name} {d.address} serial={d.serial}", err=True)
159
+ raise typer.Exit(1)
160
+ target = matches[0].address
161
+ else:
162
+ typer.echo("")
163
+ typer.echo("Found:")
164
+ for i, d in enumerate(devices, 1):
165
+ typer.echo(f" {i}. {d.name} ({d.address})")
166
+ default = "1" if len(devices) == 1 else None
167
+ idx = typer.prompt("Select device", default=default, type=int)
168
+ if not 1 <= idx <= len(devices):
169
+ typer.echo("Invalid selection.", err=True)
170
+ raise typer.Exit(1)
171
+ target = devices[idx - 1].address
172
+
173
+ wifi_ssid = ssid if ssid is not None else typer.prompt("Wi-Fi SSID")
174
+ wifi_password = (
175
+ password
176
+ if password is not None
177
+ else typer.prompt("Wi-Fi password", hide_input=True, default="", show_default=False)
178
+ )
179
+
180
+ def _on_error(attempt: int, exc: BaseException) -> None:
181
+ typer.echo(f" attempt {attempt} failed: {exc}", err=True)
182
+
183
+ device_serial = await retry(
184
+ lambda: provision_wifi(
185
+ target, # type: ignore[arg-type]
186
+ wifi_ssid,
187
+ wifi_password,
188
+ country,
189
+ ap=ap,
190
+ on_progress=lambda m: typer.echo(f" {m}"),
191
+ ),
192
+ attempts=retries,
193
+ on_error=_on_error,
194
+ )
195
+
196
+ if device_serial:
197
+ typer.echo(f"✓ Provisioned. Serial: {device_serial}")
198
+ else:
199
+ typer.echo("✓ Provisioned.")
200
+
201
+ asyncio.run(run())
202
+
203
+
204
+ def main() -> None:
205
+ app()
206
+
207
+
208
+ if __name__ == "__main__":
209
+ main()
@@ -0,0 +1,214 @@
1
+ # Copyright 2025-2026 Dimensional Inc.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ """LAN discovery for Unitree Go2 robots.
16
+
17
+ Go2s respond to a custom UDP-multicast probe (group 231.1.1.1, port 10131)
18
+ with a JSON payload containing their serial number and IP. This module sends
19
+ that probe and collects replies on port 10134.
20
+
21
+ Implemented from scratch (rather than reusing legion1581/unitree_webrtc_connect's
22
+ scanner) so we can pin the multicast send/recv to a specific interface — Tailscale
23
+ and other VPN tun devices install a 224.0.0.0/4 route in a separate table that
24
+ silently swallows the probe otherwise.
25
+ """
26
+
27
+ from __future__ import annotations
28
+
29
+ import asyncio
30
+ from collections.abc import AsyncIterator, Iterator
31
+ from dataclasses import dataclass
32
+ import json
33
+ import socket
34
+ import struct
35
+ import time
36
+
37
+ import logging
38
+
39
+ import psutil
40
+
41
+ logger = logging.getLogger(__name__)
42
+
43
+ MULTICAST_GROUP = "231.1.1.1"
44
+ QUERY_PORT = 10131
45
+ REPLY_PORT = 10134
46
+ QUERY_PAYLOAD = json.dumps({"name": "unitree_dapengche"}).encode("utf-8")
47
+
48
+
49
+ @dataclass(frozen=True)
50
+ class Go2Device:
51
+ serial: str
52
+ ip: str
53
+ iface: str
54
+ mac: str | None = None
55
+
56
+
57
+ def _read_arp(ip: str) -> str | None:
58
+ try:
59
+ with open("/proc/net/arp") as f:
60
+ for line in f.readlines()[1:]:
61
+ parts = line.split()
62
+ if len(parts) >= 4 and parts[0] == ip and parts[3] != "00:00:00:00:00:00":
63
+ return parts[3].upper()
64
+ except OSError:
65
+ return None
66
+ return None
67
+
68
+
69
+ def _resolve_mac(ip: str, retries: int = 3, retry_delay: float = 0.05) -> str | None:
70
+ """Resolve MAC for `ip` via /proc/net/arp; nudges the cache with a probe UDP send.
71
+
72
+ Linux-only. Returns None if the address can't be resolved (e.g. on macOS).
73
+ """
74
+ cached = _read_arp(ip)
75
+ if cached:
76
+ return cached
77
+ for _ in range(retries):
78
+ try:
79
+ s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
80
+ try:
81
+ s.sendto(b"", (ip, 9)) # discard port — kernel ARPs before send
82
+ except OSError:
83
+ pass
84
+ s.close()
85
+ except OSError:
86
+ pass
87
+ time.sleep(retry_delay)
88
+ mac = _read_arp(ip)
89
+ if mac:
90
+ return mac
91
+ return None
92
+
93
+
94
+ def _candidate_ifaces() -> Iterator[tuple[str, str]]:
95
+ """Yield (iface_name, ipv4_addr) for non-loopback, non-tun IPv4 interfaces."""
96
+ skip_prefixes = ("lo", "tailscale", "wg", "tun", "docker", "br-", "veth", "Meta")
97
+ for name, addrs in psutil.net_if_addrs().items():
98
+ if name.startswith(skip_prefixes):
99
+ continue
100
+ for a in addrs:
101
+ if a.family == socket.AF_INET and not a.address.startswith("127."):
102
+ yield name, a.address
103
+ break
104
+
105
+
106
+ def discover(timeout: float = 2.0, iface_ip: str | None = None) -> list[Go2Device]:
107
+ """Probe the LAN for Go2 robots.
108
+
109
+ Args:
110
+ timeout: seconds to wait for replies after sending the probe.
111
+ iface_ip: pin multicast to this local IPv4 address. If None, probe every
112
+ non-tunnel interfacel.
113
+ """
114
+ targets: list[tuple[str | None, str]] = (
115
+ [(None, iface_ip)] if iface_ip else list(_candidate_ifaces())
116
+ )
117
+ if not targets:
118
+ logger.warning("no usable interfaces found for Go2 discovery")
119
+ return []
120
+
121
+ found: dict[str, Go2Device] = {}
122
+ for name, ip in targets:
123
+ for dev in _probe_iface(ip, timeout):
124
+ dev = Go2Device(
125
+ serial=dev.serial, ip=dev.ip, iface=name or "?", mac=_resolve_mac(dev.ip)
126
+ )
127
+ found.setdefault(dev.serial, dev)
128
+ return list(found.values())
129
+
130
+
131
+ def _probe_iface(iface_ip: str, timeout: float) -> Iterator[Go2Device]:
132
+ sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
133
+ sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
134
+ try:
135
+ sock.bind(("", REPLY_PORT))
136
+ except OSError as e:
137
+ logger.warning(f"could not bind UDP {REPLY_PORT} on {iface_ip}: {e}")
138
+ sock.close()
139
+ return
140
+
141
+ iface_addr = socket.inet_aton(iface_ip)
142
+ sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_IF, iface_addr)
143
+ sock.setsockopt(
144
+ socket.IPPROTO_IP,
145
+ socket.IP_ADD_MEMBERSHIP,
146
+ struct.pack("4s4s", socket.inet_aton(MULTICAST_GROUP), iface_addr),
147
+ )
148
+
149
+ try:
150
+ sock.sendto(QUERY_PAYLOAD, (MULTICAST_GROUP, QUERY_PORT))
151
+ except OSError as e:
152
+ logger.warning(f"multicast send failed on {iface_ip}: {e}")
153
+ sock.close()
154
+ return
155
+
156
+ sock.settimeout(timeout)
157
+ try:
158
+ while True:
159
+ try:
160
+ data, addr = sock.recvfrom(1024)
161
+ except TimeoutError:
162
+ return
163
+ try:
164
+ msg = json.loads(data.decode("utf-8"))
165
+ except (UnicodeDecodeError, json.JSONDecodeError):
166
+ continue
167
+ sn = msg.get("sn")
168
+ if not sn:
169
+ continue
170
+ yield Go2Device(serial=sn, ip=msg.get("ip", addr[0]), iface=iface_ip)
171
+ finally:
172
+ sock.close()
173
+
174
+
175
+ async def discover_lan(
176
+ tick: float = 2.0,
177
+ timeout: float = 1.5,
178
+ iface_ip: str | None = None,
179
+ ) -> AsyncIterator[Go2Device]:
180
+ """Stream Go2 LAN discoveries. Polls multicast every `tick` seconds.
181
+
182
+ Yields each device on every poll cycle it shows up in (so consumers see the
183
+ device repeatedly while it remains alive on the network).
184
+ """
185
+ while True:
186
+ loop = asyncio.get_running_loop()
187
+ devs = await loop.run_in_executor(
188
+ None, lambda: discover(timeout=timeout, iface_ip=iface_ip)
189
+ )
190
+ for d in devs:
191
+ yield d
192
+ await asyncio.sleep(max(0.0, tick - timeout))
193
+
194
+
195
+ def main() -> None:
196
+ import argparse
197
+
198
+ parser = argparse.ArgumentParser(description="Discover Unitree Go2 robots on the LAN")
199
+ parser.add_argument("--timeout", type=float, default=2.0)
200
+ parser.add_argument(
201
+ "--iface-ip", help="pin probe to this local IPv4 (default: all non-tunnel ifaces)"
202
+ )
203
+ args = parser.parse_args()
204
+
205
+ devices = discover(timeout=args.timeout, iface_ip=args.iface_ip)
206
+ if not devices:
207
+ print("no Go2s found")
208
+ return
209
+ for d in devices:
210
+ print(f"{d.serial}\t{d.ip}\t(via {d.iface})")
211
+
212
+
213
+ if __name__ == "__main__":
214
+ main()