232 lines
7.1 KiB
Python
232 lines
7.1 KiB
Python
"""
|
||
wipe_device.py — Wipe ATOMS3 MacroPad NVS and LittleFS partitions.
|
||
|
||
Usage:
|
||
python wipe_device.py # auto-detect COM port
|
||
python wipe_device.py COM3 # use specified COM port
|
||
|
||
The script will:
|
||
1. Find the connected MacroPad automatically
|
||
2. Send {"cmd":"bootloader"} to trigger ROM download mode
|
||
3. Wait for the device to re-enumerate as USB-JTAG
|
||
4. Erase NVS (0x009000, 20 KB) — clears settings
|
||
5. Erase LittleFS (0x670000, 1.5 MB) — clears all macros
|
||
6. Reset the device back into normal firmware
|
||
"""
|
||
|
||
import glob
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
import serial
|
||
import serial.tools.list_ports
|
||
|
||
# Force UTF-8 output on Windows terminals
|
||
if sys.stdout.encoding and sys.stdout.encoding.lower() != "utf-8":
|
||
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||
|
||
ESPRESSIF_VID = 0x303A
|
||
DEVICE_ID = "ATOMS3-MACROPAD"
|
||
BAUD_RATE = 115200
|
||
|
||
# Partition layout (matches default_8MB.csv flashed with firmware)
|
||
NVS_OFFSET = 0x009000
|
||
NVS_SIZE = 0x005000 # 20 KB
|
||
LITTLEFS_OFFSET = 0x670000
|
||
LITTLEFS_SIZE = 0x180000 # 1.5 MB
|
||
|
||
|
||
def find_esptool() -> str:
|
||
"""Locate esptool.exe inside the Arduino15 package cache."""
|
||
pattern = os.path.join(
|
||
os.environ.get("LOCALAPPDATA", ""),
|
||
"Arduino15", "packages", "*", "tools", "esptool_py", "*", "esptool.exe",
|
||
)
|
||
matches = sorted(glob.glob(pattern), reverse=True) # newest version first
|
||
if matches:
|
||
return matches[0]
|
||
raise FileNotFoundError(
|
||
"esptool.exe not found in Arduino15 package cache.\n"
|
||
"Make sure the M5Stack board package is installed via arduino-cli."
|
||
)
|
||
|
||
|
||
def list_espressif_ports() -> set[str]:
|
||
"""Return the set of COM port names currently belonging to Espressif devices."""
|
||
return {
|
||
p.device
|
||
for p in serial.tools.list_ports.comports()
|
||
if p.vid == ESPRESSIF_VID
|
||
}
|
||
|
||
|
||
def find_macropad_port() -> str | None:
|
||
"""Scan Espressif COM ports and return the one that responds as the MacroPad."""
|
||
for p in serial.tools.list_ports.comports():
|
||
if p.vid != ESPRESSIF_VID:
|
||
continue
|
||
try:
|
||
ser = serial.Serial()
|
||
ser.port = p.device
|
||
ser.baudrate = BAUD_RATE
|
||
ser.timeout = 2
|
||
ser.dtr = False
|
||
ser.rts = False
|
||
ser.open()
|
||
time.sleep(0.1)
|
||
ser.dtr = True # signal host-connected to TinyUSB CDC
|
||
time.sleep(0.3)
|
||
ser.reset_input_buffer()
|
||
ser.write(b'{"cmd":"ping"}\n')
|
||
line = ser.readline().decode("utf-8", errors="ignore").strip()
|
||
ser.close()
|
||
if line:
|
||
data = json.loads(line)
|
||
if data.get("id") == DEVICE_ID:
|
||
return p.device
|
||
except (serial.SerialException, json.JSONDecodeError, OSError):
|
||
pass
|
||
return None
|
||
|
||
|
||
def send_bootloader_command(port: str) -> None:
|
||
"""Open the port and fire the bootloader command. Ignores read errors on disconnect."""
|
||
ser = serial.Serial()
|
||
ser.port = port
|
||
ser.baudrate = BAUD_RATE
|
||
ser.timeout = 2
|
||
ser.dtr = False
|
||
ser.rts = False
|
||
ser.open()
|
||
time.sleep(0.1)
|
||
ser.dtr = True
|
||
time.sleep(0.3)
|
||
ser.reset_input_buffer()
|
||
ser.write(b'{"cmd":"bootloader"}\n')
|
||
ser.flush()
|
||
try:
|
||
# Read {"rsp":"ok"} — may throw if device resets mid-read
|
||
ser.read(64)
|
||
except (serial.SerialException, OSError):
|
||
pass
|
||
finally:
|
||
try:
|
||
ser.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def wait_for_port_change(old_ports: set[str], timeout: float = 10.0) -> str:
|
||
"""
|
||
Block until a new Espressif COM port appears that wasn't in old_ports.
|
||
Returns the new port name, or raises TimeoutError.
|
||
"""
|
||
deadline = time.time() + timeout
|
||
while time.time() < deadline:
|
||
current = list_espressif_ports()
|
||
new = current - old_ports
|
||
if new:
|
||
return new.pop()
|
||
time.sleep(0.25)
|
||
raise TimeoutError(
|
||
"Timed out waiting for device to re-enumerate in download mode.\n"
|
||
"Check USB connection and try again."
|
||
)
|
||
|
||
|
||
def erase_region(esptool: str, port: str, offset: int, size: int, keep_stub: bool) -> None:
|
||
"""Run esptool to erase one flash region.
|
||
|
||
keep_stub=True leaves the stub running so the next call can reconnect
|
||
without a full reset cycle.
|
||
"""
|
||
after = "no-reset" if keep_stub else "watchdog-reset"
|
||
cmd = [
|
||
esptool,
|
||
"--chip", "esp32s3",
|
||
"--port", port,
|
||
"--baud", "460800",
|
||
"--before", "no-reset",
|
||
"--after", after,
|
||
"erase-region",
|
||
hex(offset),
|
||
hex(size),
|
||
]
|
||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||
output = result.stdout + result.stderr
|
||
if result.returncode != 0:
|
||
print(output)
|
||
raise RuntimeError(f"esptool failed (exit {result.returncode})")
|
||
|
||
|
||
def main() -> None:
|
||
print("=" * 40)
|
||
print(" ATOMS3 MacroPad -- Device Wiper")
|
||
print("=" * 40 + "\n")
|
||
|
||
print("[ 1/5 ] Locating esptool...", end=" ", flush=True)
|
||
esptool = find_esptool()
|
||
print(f"OK\n {esptool}")
|
||
|
||
forced_port = sys.argv[1].upper() if len(sys.argv) > 1 else None
|
||
if forced_port:
|
||
print(f"[ 2/5 ] Using specified port {forced_port}...", end=" ", flush=True)
|
||
port = forced_port
|
||
print("OK")
|
||
else:
|
||
print("[ 2/5 ] Searching for MacroPad...", end=" ", flush=True)
|
||
port = find_macropad_port()
|
||
if not port:
|
||
print("NOT FOUND")
|
||
print("\nNo MacroPad detected. Make sure the device is connected and firmware is running.")
|
||
sys.exit(1)
|
||
print(f"found on {port}")
|
||
|
||
print("[ 3/5 ] Entering download mode...", end=" ", flush=True)
|
||
ports_before = list_espressif_ports()
|
||
send_bootloader_command(port)
|
||
# Wait for old port to vanish, then for the JTAG port to appear
|
||
time.sleep(0.5)
|
||
try:
|
||
jtag_port = wait_for_port_change(ports_before - {port}, timeout=10.0)
|
||
except TimeoutError as e:
|
||
print("TIMEOUT")
|
||
print(f"\n{e}")
|
||
sys.exit(1)
|
||
print(f"device on {jtag_port}")
|
||
|
||
# Let the ROM fully settle
|
||
time.sleep(0.5)
|
||
|
||
print(f"[ 4/5 ] Erasing NVS "
|
||
f"(0x{NVS_OFFSET:06X} – 0x{NVS_OFFSET + NVS_SIZE - 1:06X}, "
|
||
f"{NVS_SIZE // 1024} KB)...", end=" ", flush=True)
|
||
try:
|
||
erase_region(esptool, jtag_port, NVS_OFFSET, NVS_SIZE, keep_stub=True)
|
||
except RuntimeError as e:
|
||
print("FAILED")
|
||
print(f"\n{e}")
|
||
sys.exit(1)
|
||
print("OK")
|
||
|
||
# Final erase resets the device
|
||
print(f"[ 5/5 ] Erasing LittleFS "
|
||
f"(0x{LITTLEFS_OFFSET:06X} – 0x{LITTLEFS_OFFSET + LITTLEFS_SIZE - 1:06X}, "
|
||
f"{LITTLEFS_SIZE // 1024} KB)...", end=" ", flush=True)
|
||
try:
|
||
erase_region(esptool, jtag_port, LITTLEFS_OFFSET, LITTLEFS_SIZE, keep_stub=False)
|
||
except RuntimeError as e:
|
||
print("FAILED")
|
||
print(f"\n{e}")
|
||
sys.exit(1)
|
||
print("OK")
|
||
|
||
print("\nDone. NVS and LittleFS are blank -- device is in a fresh state.\n")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|