Files
2026-07-17 15:29:53 -04:00

445 lines
17 KiB
Python

"""Serial communication manager for ATOMS3 MacroPad."""
import json
import os
import time
import threading
import serial
import serial.tools.list_ports
import ble_keystore
from utils.constants import DEVICE_ID, ESPRESSIF_VID
from utils.image_converter import convert_to_rgb565
class SerialManager:
def __init__(self):
self.ser: serial.Serial | None = None
self.port: str | None = None
self.connected = False
self.is_uploading = False
# Last ping response (board, sta_mac, mesh_ch, live_tx, ...). Captured
# on connect and refreshed on every ping() so the UI can show which
# device is attached and what live transport it's set to.
self.device_info: dict | None = None
self._lock = threading.Lock()
self._on_disconnect = None
self._on_connect = None
def set_callbacks(self, on_connect=None, on_disconnect=None):
self._on_connect = on_connect
self._on_disconnect = on_disconnect
def scan_and_connect(self) -> bool:
"""Scan COM ports for ATOMS3 MacroPad device."""
ports = serial.tools.list_ports.comports()
for port in ports:
if port.vid == ESPRESSIF_VID:
try:
# Open without DTR/RTS asserted so the ESP32 doesn't reset
ser = serial.Serial()
ser.port = port.device
ser.baudrate = 115200
ser.timeout = 2
ser.dtr = False
ser.rts = False
ser.open()
time.sleep(0.1)
# Assert DTR so TinyUSB CDC sees the host as connected
ser.dtr = True
time.sleep(0.3)
ser.reset_input_buffer()
ser.write(b'{"cmd":"ping"}\n')
line = ser.readline().decode("utf-8", errors="ignore").strip()
if line:
data = json.loads(line)
if data.get("id") == DEVICE_ID:
self.ser = ser
self.port = port.device
self.connected = True
self.device_info = data
if self._on_connect:
self._on_connect(port.device)
return True
ser.close()
except (serial.SerialException, json.JSONDecodeError, OSError):
pass
return False
def disconnect(self):
if self.ser and self.ser.is_open:
self.ser.close()
self.ser = None
self.port = None
self.connected = False
self.device_info = None
if self._on_disconnect:
self._on_disconnect()
def send_command(self, cmd: dict) -> dict | None:
"""Send a JSON command and wait for response."""
if not self.connected or not self.ser:
return None
with self._lock:
try:
msg = json.dumps(cmd) + "\n"
self.ser.write(msg.encode("utf-8"))
line = self.ser.readline().decode("utf-8", errors="ignore").strip()
if line:
return json.loads(line)
return None
except (serial.SerialException, json.JSONDecodeError, OSError):
self.connected = False
if self._on_disconnect:
self._on_disconnect()
return None
def ping(self) -> dict | None:
"""Ping the device. Newer firmware also returns 'board'
("atoms3"/"atoms3_lite"), 'sta_mac' (the WiFi/mesh MAC), 'mesh_ch',
and 'live_tx' (0 = mesh, 1 = BLE) alongside the legacy
id/ver/macros/free fields."""
rsp = self.send_command({"cmd": "ping"})
if rsp:
self.device_info = rsp
return rsp
def live_mode_label(self) -> str | None:
"""Human label for the connected device's persisted live transport,
or None if unknown (older firmware doesn't report it)."""
if not self.device_info:
return None
tx = self.device_info.get("live_tx")
if tx is None:
return None
return "BLE" if tx == 1 else "Mesh"
def get_ble_key_info(self) -> dict | None:
"""Return {'key': bytes, 'tag': str|None} from the device.
Newer firmware includes the device tag (with MAC) so the host
can store keys per-device. Older firmware returns just the key
— tag will be None.
"""
rsp = self.send_command({"cmd": "get_ble_key"})
if rsp is None or rsp.get("rsp") != "ble_key":
return None
try:
key = bytes.fromhex(rsp["key"])
except (KeyError, ValueError):
return None
if len(key) != ble_keystore.KEY_LEN:
return None
tag = rsp.get("tag")
if not isinstance(tag, str):
tag = None
return {"key": key, "tag": tag}
def get_ble_key(self) -> bytes | None:
"""Pull the firmware's BLE payload-encryption key (32 bytes)."""
rsp = self.send_command({"cmd": "get_ble_key"})
if rsp and rsp.get("rsp") == "ble_key":
try:
key = bytes.fromhex(rsp["key"])
except (KeyError, ValueError):
return None
return key if len(key) == ble_keystore.KEY_LEN else None
return None
def get_settings(self) -> dict | None:
return self.send_command({"cmd": "get_settings"})
def set_setting(self, key: str, value: int) -> bool:
rsp = self.send_command({"cmd": "set", "key": key, "val": value})
return rsp is not None and rsp.get("rsp") == "ok"
def upload_settings(self, settings) -> bool:
"""Upload all settings to device."""
ok = True
ok = ok and self.set_setting("hold_ms", settings.hold_ms)
ok = ok and self.set_setting("orientation", settings.orientation)
ok = ok and self.set_setting("type_delay", settings.type_delay)
ok = ok and self.set_setting("resume_delay", settings.resume_delay)
ok = ok and self.set_setting("combo_pre_ms", settings.combo_pre_ms)
ok = ok and self.set_setting("combo_post_ms", settings.combo_post_ms)
ok = ok and self.set_setting("probe_timeout_ms", settings.probe_timeout_ms)
ok = ok and self.set_setting("media_hold_ms", settings.media_hold_ms)
ok = ok and self.set_setting("type_shift_extra_ms", settings.type_shift_extra_ms)
ok = ok and self.set_setting("type_settle_ms", settings.type_settle_ms)
ok = ok and self.set_setting("pause_margin_left", settings.pause_margin_left)
ok = ok and self.set_setting("pause_margin_right", settings.pause_margin_right)
ok = ok and self.set_setting("pause_margin_top", settings.pause_margin_top)
ok = ok and self.set_setting("pause_margin_bottom", settings.pause_margin_bottom)
return ok
def upload_macro(self, slot: int, macro, progress_cb=None) -> bool:
"""Upload a single macro to the device."""
if not self.connected or not self.ser:
return False
with self._lock:
try:
nodes = macro.flatten_for_device()
node_count = len(nodes)
img_data = b""
if macro.image_path and os.path.exists(macro.image_path):
img_data = convert_to_rgb565(macro.image_path)
begin_cmd = json.dumps({
"cmd": "macro_begin",
"slot": slot,
"name": macro.name,
"label_color": getattr(macro, "label_color", "white"),
"node_count": node_count,
"img_size": len(img_data),
}) + "\n"
self.ser.write(begin_cmd.encode("utf-8"))
rsp = self._read_response()
if not rsp or rsp.get("rsp") != "ready":
return False
if progress_cb:
progress_cb(0.1)
if img_data:
chunk_size = 128 # Small chunks for reliable TinyUSB CDC transfer
total_chunks = (len(img_data) + chunk_size - 1) // chunk_size
for ci, i in enumerate(range(0, len(img_data), chunk_size)):
chunk = img_data[i:i + chunk_size]
self.ser.write(f"CHUNK:{len(chunk)}\n".encode("utf-8"))
self.ser.flush()
self.ser.write(chunk)
self.ser.flush()
ack = self.ser.readline().decode("utf-8", errors="ignore").strip()
if ack != "OK":
return False
if progress_cb:
progress_cb(0.1 + 0.3 * (ci + 1) / total_chunks)
self.ser.write(b"IMG_DONE\n")
self.ser.flush()
rsp = self._read_response()
if not rsp or rsp.get("rsp") != "img_ok":
return False
# Device needs a moment to settle after the image write
time.sleep(0.2)
self.ser.reset_input_buffer()
if progress_cb:
progress_cb(0.4)
for i, node in enumerate(nodes):
node_cmd = json.dumps({
"cmd": "node",
"idx": i,
"type": node["type"],
"data": node["data"],
}) + "\n"
self.ser.write(node_cmd.encode("utf-8"))
rsp = self._read_response()
if not rsp or rsp.get("rsp") != "ok":
return False
if progress_cb and node_count > 0:
progress_cb(0.4 + 0.4 * (i + 1) / node_count)
end_cmd = json.dumps({"cmd": "macro_end", "slot": slot}) + "\n"
self.ser.write(end_cmd.encode("utf-8"))
rsp = self._read_response()
if not rsp or rsp.get("rsp") != "ok":
return False
if progress_cb:
progress_cb(1.0)
return True
except (serial.SerialException, OSError):
self.connected = False
if self._on_disconnect:
self._on_disconnect()
return False
def upload_all(self, project, progress_cb=None, subroutine_macros=None) -> bool:
"""Upload entire project (settings + all macros + subroutines) to device.
Deletes all existing macros on the device first, then re-uploads everything.
"""
self.is_uploading = True
try:
return self._upload_all_impl(project, progress_cb, subroutine_macros)
finally:
self.is_uploading = False
def _upload_all_impl(self, project, progress_cb=None, subroutine_macros=None) -> bool:
if not self.connected:
return False
ping_rsp = self.ping()
if not ping_rsp:
return False
# Sync the BLE payload-encryption key from firmware. Older firmware
# without this command returns None — leave any cached key alone.
# Newer firmware also returns the device tag (with MAC) so the
# host can keep keys for multiple ATOMS3s without clobbering each
# other on every upload.
info = self.get_ble_key_info()
if info is not None:
try:
# Always update the legacy single-key file (back-compat).
ble_keystore.save_key(info["key"])
# And the per-MAC store, if the device reported its tag.
if info["tag"]:
ble_keystore.save_key_for_mac(info["tag"], info["key"])
except OSError as exc:
print(f"[BLE] Failed to persist key: {exc}")
device_macro_count = ping_rsp.get("macros", 0)
for slot in range(device_macro_count):
self.delete_macro(slot)
time.sleep(0.05)
self.send_command({"cmd": "sub_clear"})
time.sleep(0.05)
sub_macros = subroutine_macros or []
total = len(project.macros) + len(sub_macros) + 2 # +2: delete phase, settings
step = 1
if progress_cb:
progress_cb(step / total)
if not self.upload_settings(project.settings):
return False
step += 1
if progress_cb:
progress_cb(step / total)
for i, macro in enumerate(project.macros):
def macro_progress(p, _step=step):
if progress_cb:
progress_cb((_step + p) / total)
if not self.upload_macro(i, macro, macro_progress):
return False
step += 1
if progress_cb:
progress_cb(step / total)
order = list(range(len(project.macros)))
self.send_command({"cmd": "macro_reorder", "order": order})
for i, sub_macro in enumerate(sub_macros):
nodes = sub_macro.flatten_for_device()
node_count = len(nodes)
begin_rsp = self.send_command({
"cmd": "sub_begin",
"slot": i,
"name": sub_macro.name,
"node_count": node_count,
})
if not begin_rsp or begin_rsp.get("rsp") != "ready":
# Older firmware without sub-routine support replies non-ready
step += 1
if progress_cb:
progress_cb(step / total)
continue
node_ok = True
for j, node in enumerate(nodes):
rsp = self.send_command({
"cmd": "sub_node",
"idx": j,
"type": node["type"],
"data": node["data"],
})
if not rsp or rsp.get("rsp") != "ok":
node_ok = False
break
end_rsp = self.send_command({"cmd": "sub_end", "slot": i})
if not node_ok or not end_rsp or end_rsp.get("rsp") != "ok":
return False
step += 1
if progress_cb:
progress_cb(step / total)
return True
def delete_macro(self, slot: int) -> bool:
rsp = self.send_command({"cmd": "macro_delete", "slot": slot})
return rsp is not None and rsp.get("rsp") == "ok"
def rs232_open(self, baud: int, data_bits: int = 8,
stop_bits: str = "1", parity: str = "none") -> bool:
rsp = self.send_command({
"cmd": "rs232_open",
"baud": baud,
"data_bits": data_bits,
"stop_bits": stop_bits,
"parity": parity,
})
return rsp is not None and rsp.get("rsp") == "ok"
def rs232_close(self) -> bool:
rsp = self.send_command({"cmd": "rs232_close"})
return rsp is not None and rsp.get("rsp") == "ok"
def rs232_send(self, data: bytes) -> bool:
"""Send raw bytes over the RS232 port (host-side terminal)."""
if not isinstance(data, (bytes, bytearray)):
data = str(data).encode("utf-8", errors="replace")
hex_str = data.hex()
rsp = self.send_command({"cmd": "rs232_send", "hex": hex_str})
return rsp is not None and rsp.get("rsp") == "ok"
def rs232_poll(self) -> bytes:
"""Return any bytes received since the last poll (may be empty)."""
rsp = self.send_command({"cmd": "rs232_poll"})
if rsp and rsp.get("rsp") == "rx":
hex_str = rsp.get("hex", "")
if hex_str:
try:
return bytes.fromhex(hex_str)
except ValueError:
return b""
return b""
def get_log(self) -> list | None:
rsp = self.send_command({"cmd": "get_log"})
if rsp and rsp.get("rsp") == "log":
return rsp.get("entries", [])
return None
def clear_log(self) -> bool:
rsp = self.send_command({"cmd": "clear_log"})
return rsp is not None and rsp.get("rsp") == "ok"
def get_ble_log(self) -> list | None:
"""Pull the device's in-memory BLE debug ring buffer."""
rsp = self.send_command({"cmd": "get_ble_log"})
if rsp and rsp.get("rsp") == "ble_log":
return rsp.get("entries", [])
return None
def clear_ble_log(self) -> bool:
rsp = self.send_command({"cmd": "clear_ble_log"})
return rsp is not None and rsp.get("rsp") == "ok"
def _read_response(self, timeout: float = 10.0) -> dict | None:
"""Read a JSON response line with timeout."""
old_timeout = self.ser.timeout
self.ser.timeout = timeout
try:
line = self.ser.readline().decode("utf-8", errors="ignore").strip()
if line:
return json.loads(line)
return None
except (json.JSONDecodeError, serial.SerialException):
return None
finally:
self.ser.timeout = old_timeout