204 lines
7.6 KiB
Python
204 lines
7.6 KiB
Python
|
||
#!/usr/bin/env python3
|
||
import argparse, os, sys, time, fnmatch
|
||
from pathlib import Path
|
||
from typing import List, Dict
|
||
|
||
import serial, serial.tools.list_ports as list_ports
|
||
import subprocess
|
||
|
||
KNOWN_VIDS = {0x0403, 0x10C4, 0x1A86, 0x067B, 0x303A}
|
||
|
||
def run(cmd, cwd=None, check=True):
|
||
print(">>", " ".join(cmd))
|
||
p = subprocess.run(cmd, cwd=cwd)
|
||
if check and p.returncode != 0:
|
||
raise RuntimeError(f"Command failed ({p.returncode}): {' '.join(cmd)}")
|
||
return p.returncode
|
||
|
||
def detect_chip_type(port: str) -> str:
|
||
try:
|
||
out = subprocess.check_output(["esptool.py", "--port", port, "chip_id"], stderr=subprocess.STDOUT, text=True, timeout=6)
|
||
if "ESP32-S3" in out.upper(): return "ESP32-S3"
|
||
if "ESP32-S2" in out.upper(): return "ESP32-S2"
|
||
if "ESP32-C3" in out.upper(): return "ESP32-C3"
|
||
if "ESP32-C6" in out.upper(): return "ESP32-C6"
|
||
if "ESP32" in out.upper(): return "ESP32"
|
||
except Exception:
|
||
pass
|
||
return "Unknown"
|
||
|
||
def map_chip_to_idf_target(chip: str) -> str:
|
||
s = chip.upper()
|
||
if s.startswith("ESP32-S3"): return "esp32s3"
|
||
if s.startswith("ESP32-S2"): return "esp32s2"
|
||
if s.startswith("ESP32-C3"): return "esp32c3"
|
||
if s.startswith("ESP32-C6"): return "esp32c6"
|
||
if s.startswith("ESP32-H2"): return "esp32h2"
|
||
if s.startswith("ESP32"): return "esp32"
|
||
return "unknown"
|
||
|
||
def list_ports_filtered(patterns: List[str] = None):
|
||
ports = list_ports.comports()
|
||
out = []
|
||
for p in ports:
|
||
dev = p.device
|
||
if patterns:
|
||
if not any(fnmatch.fnmatch(dev, pat) for pat in patterns):
|
||
continue
|
||
else:
|
||
if not (dev.startswith("/dev/ttyUSB") or dev.startswith("/dev/ttyACM")):
|
||
continue
|
||
vid = getattr(p, "vid", None)
|
||
if vid is not None and vid not in KNOWN_VIDS:
|
||
continue
|
||
out.append(p)
|
||
return out
|
||
|
||
def ensure_target(project_dir: str, target: str):
|
||
if not target or target == "unknown":
|
||
raise ValueError("Unknown IDF target; cannot set-target")
|
||
run(["idf.py", "set-target", target], cwd=project_dir, check=True)
|
||
|
||
def flash_device(project_dir: str, port: str, idf_target: str, baud: int = 460800) -> bool:
|
||
try:
|
||
ensure_target(project_dir, idf_target)
|
||
run(["idf.py", "-p", port, "-b", str(baud), "flash"], cwd=project_dir, check=True)
|
||
return True
|
||
except Exception as e:
|
||
print(f" Flash failed on {port}: {e}")
|
||
return False
|
||
|
||
def toggle_reset(ser):
|
||
try:
|
||
ser.dtr = False
|
||
ser.rts = True
|
||
time.sleep(0.05)
|
||
ser.dtr = True
|
||
ser.rts = False
|
||
time.sleep(0.05)
|
||
except Exception:
|
||
pass
|
||
|
||
def send_wifi_config(port: str, ssid: str, password: str, ip: str, mask: str, gw: str, dhcp: bool, baud: int = 115200, retries: int = 3) -> bool:
|
||
for attempt in range(1, retries+1):
|
||
try:
|
||
print(f" [{port}] opening serial @ {baud} (attempt {attempt}/{retries})")
|
||
with serial.Serial(port, baudrate=baud, timeout=2) as ser:
|
||
time.sleep(0.3)
|
||
toggle_reset(ser)
|
||
time.sleep(0.8)
|
||
ser.reset_input_buffer()
|
||
ser.reset_output_buffer()
|
||
|
||
lines = [
|
||
"CFG\n",
|
||
f"SSID:{ssid}\n" if ssid else "",
|
||
f"PASS:{password}\n" if password else "",
|
||
f"IP:{ip}\n" if ip else "",
|
||
f"MASK:{mask}\n" if mask else "",
|
||
f"GW:{gw}\n" if gw else "",
|
||
f"DHCP:{1 if dhcp else 0}\n",
|
||
"END\n",
|
||
]
|
||
payload = "".join([l for l in lines if l])
|
||
ser.write(payload.encode("utf-8"))
|
||
ser.flush()
|
||
|
||
t0 = time.time()
|
||
buf = b""
|
||
while time.time() - t0 < 3.0:
|
||
chunk = ser.read(64)
|
||
if chunk:
|
||
buf += chunk
|
||
if b"OK" in buf:
|
||
print(f" [{port}] config applied: {buf.decode(errors='ignore').strip()}")
|
||
return True
|
||
print(f" [{port}] no OK from device, got: {buf.decode(errors='ignore').strip()}")
|
||
except Exception as e:
|
||
print(f" [{port}] serial error: {e}")
|
||
time.sleep(0.6)
|
||
return False
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description="Mass flash ESP32 devices and push Wi‑Fi/IP config over serial.")
|
||
ap.add_argument("--project", required=True, help="Path to ESP‑IDF project")
|
||
ap.add_argument("--ssid", required=True, help="Wi‑Fi SSID")
|
||
ap.add_argument("--password", required=True, help="Wi‑Fi password")
|
||
ap.add_argument("--start-ip", default="192.168.1.50", help="Base IP (only used if --dhcp=0)")
|
||
ap.add_argument("--mask", default="255.255.255.0", help="Netmask for static IP")
|
||
ap.add_argument("--gw", default="192.168.1.1", help="Gateway for static IP")
|
||
ap.add_argument("--dhcp", type=int, choices=[0,1], default=1, help="1=use DHCP, 0=set static IPs")
|
||
ap.add_argument("--baud", type=int, default=460800, help="Flashing baud rate")
|
||
ap.add_argument("--cfg-baud", type=int, default=115200, help="Serial baud for config exchange")
|
||
ap.add_argument("--ports", help="Comma-separated globs to override port selection, e.g. '/dev/ttyUSB*,/dev/ttyACM*'")
|
||
ap.add_argument("--dry-run", action="store_true", help="Plan only; do not flash or configure")
|
||
args = ap.parse_args()
|
||
|
||
project_dir = os.path.abspath(args.project)
|
||
if not os.path.isdir(project_dir):
|
||
print(f"Project directory not found: {project_dir}")
|
||
sys.exit(1)
|
||
|
||
patterns = [p.strip() for p in args.ports.split(",")] if args.ports else None
|
||
devices = list_ports_filtered(patterns)
|
||
if not devices:
|
||
print("No candidate USB serial ports found.")
|
||
sys.exit(2)
|
||
print(f"Found {len(devices)} devices.")
|
||
|
||
base = args.start_ip.split(".")
|
||
try:
|
||
base = [int(x) for x in base]
|
||
except Exception:
|
||
base = [192,168,1,50]
|
||
|
||
plan = []
|
||
for idx, dev in enumerate(devices, 1):
|
||
port = dev.device
|
||
chip = detect_chip_type(port)
|
||
target = map_chip_to_idf_target(chip)
|
||
ip = f"{base[0]}.{base[1]}.{base[2]}.{base[3] + idx - 1}" if args.dhcp == 0 else None
|
||
plan.append(dict(idx=idx, port=port, chip=chip, target=target, ip=ip))
|
||
|
||
print("\nPlan:")
|
||
for d in plan:
|
||
print(f" {d['idx']:2d} {d['port']} {d['chip']} -> target {d['target']} IP:{d['ip'] or 'DHCP'}")
|
||
|
||
if args.dry_run:
|
||
print("\nDry run only.")
|
||
return
|
||
|
||
failed = []
|
||
for d in plan:
|
||
if d['target'] == 'unknown':
|
||
print(f"\nERROR: Unknown IDF target for {d['port']} (chip '{d['chip']}'). Skipping.")
|
||
failed.append(d['idx'])
|
||
continue
|
||
print(f"\nFlashing {d['port']} as {d['target']}...")
|
||
if not flash_device(project_dir, d['port'], d['target'], baud=args.baud):
|
||
failed.append(d['idx'])
|
||
continue
|
||
|
||
print(f"Configuring Wi‑Fi on {d['port']}...")
|
||
ok = send_wifi_config(
|
||
d['port'],
|
||
args.ssid,
|
||
args.password,
|
||
d['ip'],
|
||
args.mask if d['ip'] else None,
|
||
args.gw if d['ip'] else None,
|
||
dhcp=(args.dhcp == 1),
|
||
baud=args.cfg_baud,
|
||
)
|
||
if not ok:
|
||
print(f" WARN: config not acknowledged on {d['port']}")
|
||
|
||
if failed:
|
||
print(f"\nCompleted with flashing failures on: {failed}")
|
||
sys.exit(3)
|
||
print("\nAll done.")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|