EnvelopeCodec message layer; scripts are thin business-API wrappers

franklinwh_agate_client/battery.py

@@ -258,12 +258,9 @@ def _recv_all(sock: socket.socket, idle: float = 4.0) -> bytes:
 
 def _apower_tx(sn: str, key_radix: int) -> Tuple[bytes, Dict[str, Any]]:
     """Canonical ``queryAPowerData`` wire bytes (fresh snno/timeStamp)."""
-    from .protocol import (
-        direct_connect_request,
-        encrypt_object,
-        generate_secret_key,
-    )
+    from .protocol import direct_connect_request
     from .protocol.envelope import SequenceNumber
+    from .protocol.messages import EnvelopeCodec
 
     _apower_tx_seq = getattr(_apower_tx, "_seq", None)
     if _apower_tx_seq is None:
@@ -272,9 +269,8 @@ def _apower_tx(sn: str, key_radix: int) -> Tuple[bytes, Dict[str, Any]]:
         CommunicationCmd.DEVICE_APOWER_INFO, sn,
         data={"update": 2}, snno=_apower_tx_seq.next(),
     )
-    wire = encrypt_object(
-        env.to_json_string(), generate_secret_key(sn), radix=key_radix)
-    return wire, env.to_dict()
+    codec = EnvelopeCodec.for_serial(sn, radix=key_radix)
+    return codec.encode(env), env.to_dict()
 
 
 #: Minimal MQTT v3.1.1 CONNECT (client id ``probe01``) — rules out / confirms
@@ -297,7 +293,7 @@ def probe_transport(
     prefixes (LE16/BE16/LE32), an MQTT CONNECT, and a double-send over one
     persistent connection.  Each entry carries exact ``tx``/``rx`` bytes.
     """
-    from .protocol import decrypt_wire
+    from .protocol.messages import EnvelopeCodec
 
     results: List[Dict[str, Any]] = []
     for port in ports:
@@ -330,15 +326,8 @@ def probe_transport(
             except (ConnectionError, socket.error, TimeoutError) as exc:
                 entry["error"] = f"{type(exc).__name__}: {exc}"
             if entry["rx"]:
-                try:
-                    decoded = decrypt_wire(entry["rx"], radix=key_radix)
-                except ValueError:
-                    decoded = ""
-                if decoded:
-                    try:
-                        entry["reply"] = json.loads(decoded)
-                    except ValueError:
-                        entry["reply"] = {"_ undecodable": decoded}
+                entry["reply"] = EnvelopeCodec.for_serial(
+                    sn, radix=key_radix).try_decode(entry["rx"])
             results.append(entry)
 
         run("greet-first/read-then-raw", tx, pre_read=4.0)

franklinwh_agate_client/device_local.py

@@ -19,7 +19,6 @@ the direct-connect path of the app (see ``docs/RE_DART_FINDINGS.md`` §8,
 
 from __future__ import annotations
 
-import json
 import logging
 import socket
 from typing import Any, Dict, Optional
@@ -32,12 +31,11 @@ from .protocol import (
     CommunicationOperate,
     SequenceNumber,
     SocketFrameDecoder,
-    decrypt_wire,
     direct_connect_request,
-    encrypt_object,
     frame_socket_payload,
     generate_secret_key,
 )
+from .protocol.messages import EnvelopeCodec
 from .wifi import DEVICE_GATEWAY_IP
 
 log = logging.getLogger(__name__)
@@ -96,6 +94,7 @@ class LocalDeviceClient:
         self.key_radix = int(key_radix)
         self.timeout = timeout
         self.key = generate_secret_key(sn)
+        self.codec = EnvelopeCodec(self.key, radix=self.key_radix)
         self._sn_seq = SequenceNumber()
         self._session = session or requests.Session()
         self._session.headers.update(
@@ -196,57 +195,60 @@ class LocalDeviceClient:
         return self.send_envelope(envelope)
 
     def send_envelope(self, envelope: CommunicationObject) -> Dict[str, Any]:
-        wire = encrypt_object(envelope.to_json_string(), self.key,
-                              radix=self.key_radix)
-        raw = self._socket_roundtrip(frame_socket_payload(wire))
+        raw = self._socket_roundtrip(
+            frame_socket_payload(self.codec.encode(envelope)))
         try:
-            reply = decrypt_wire(raw, radix=self.key_radix)
+            reply = self.codec.decode(raw).to_dict()
         except ValueError:
             raise ConnectionError(
                 f"undecodable reply ({len(raw)} bytes) from "
                 f"{self.socket_host}:{self.socket_port}"
             )
-        return json.loads(reply)
+        return reply
 
     def send_probe(
         self,
         cmd: CommunicationCmd,
         data: Optional[Dict[str, Any]] = None,
         timeout: Optional[float] = None,
+        equip_no: Optional[str] = None,
     ) -> Dict[str, Any]:
         """Send one envelope, return exact TX/RX bytes plus any parsed reply.
 
-        Never raises on a device-level failure: ``reply`` is the parsed JSON
-        (or ``None``), ``reply_raw`` the exact bytes received (``b""`` when
-        the device accepted the frame but answered with zero bytes).
+        Never raises on a device-level failure: ``reply`` is the parsed
+        envelope dict (or ``None``), ``reply_raw`` the exact bytes received
+        (``b""`` when the device accepted the frame but answered with zero
+        bytes).  ``equip_no`` overrides the serial (``""`` reproduces the
+        app's empty-SN discovery frame).
         """
         envelope = direct_connect_request(
-            cmd, self.sn, data=data, snno=self._sn_seq.next(),
+            cmd, self.sn if equip_no is None else equip_no, data=data,
+            snno=self._sn_seq.next(),
         )
-        wire = encrypt_object(envelope.to_json_string(), self.key,
-                              radix=self.key_radix)
+        wire = self.codec.encode(envelope)
         reply_raw = self._socket_roundtrip(
             frame_socket_payload(wire), timeout=timeout, collect_only=True,
         )
-        reply: Optional[Dict[str, Any]] = None
-        if reply_raw:
-            try:
-                decoded = decrypt_wire(reply_raw, radix=self.key_radix)
-            except ValueError:
-                decoded = ""
-            if decoded:
-                try:
-                    reply = json.loads(decoded)
-                except ValueError:
-                    reply = None
         return {
             "cmd": str(cmd),
             "envelope": envelope.to_dict(),
             "tx": wire,
             "reply_raw": reply_raw,
-            "reply": reply,
+            "reply": self.codec.try_decode(reply_raw),
         }
 
+    def send_raw(self, payload: bytes,
+                 timeout: Optional[float] = None) -> bytes:
+        """Low-level escape hatch: one connection, send *payload*, drain.
+
+        Used by transport-level probes (framing/handshake variants) that
+        are not valid envelopes.  Returns the exact bytes received (``b""``
+        when the peer answers with nothing).  Envelope traffic must use
+        :meth:`send_envelope` / :meth:`send_probe` instead.
+        """
+        return self._socket_roundtrip(bytes(payload), timeout=timeout,
+                                      collect_only=True)
+
     def _socket_roundtrip(self, framed: bytes,
                           timeout: Optional[float] = None,
                           collect_only: bool = False) -> bytes:
@@ -278,7 +280,7 @@ class LocalDeviceClient:
         if collect_only:
             return bytes(received)
         try:
-            decrypt_wire(bytes(received), radix=self.key_radix)
+            self.codec.decode(bytes(received))
         except ValueError:
             raise ConnectionError(
                 f"no reply frame from {self.socket_host}:{self.socket_port}"

franklinwh_agate_client/protocol/__init__.py

@@ -61,6 +61,7 @@ from .framing import (
     verify_frame_crc,
 )
 from .framing import SocketFrameDecoder
+from .messages import EnvelopeCodec
 
 __all__ = [
     "CommunicationCmd",
@@ -106,4 +107,5 @@ __all__ = [
     "chunk_for_mtu",
     "frame_socket_payload",
     "SocketFrameDecoder",
+    "EnvelopeCodec",
 ]

franklinwh_agate_client/protocol/messages.py

@@ -0,0 +1,64 @@
+"""Low-level wire messages: envelopes to bytes and back.
+
+:class:`EnvelopeCodec` owns everything byte-level about one device key:
+serialising a :class:`CommunicationObject` to encrypted TX bytes and
+parsing received RX bytes back into an envelope.  Business logic (which
+command/arguments to send, retries, SoC extraction) lives one layer up in
+``device_local`` / ``battery`` / ``power``; standalone scripts must go
+through those layers, never reimplement this.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Dict
+
+from .crypto import (
+    DEFAULT_KEY_RADIX,
+    decrypt_wire,
+    encrypt_object,
+    generate_secret_key,
+)
+from .envelope import CommunicationObject
+
+
+class EnvelopeCodec:
+    """Encrypt/decrypt envelopes under one device key.
+
+    :param key: ``generate_secret_key`` output for the device serial.
+    :param radix: key-parse radix (16 — verified against live captures).
+    """
+
+    def __init__(self, key: str, radix: int = DEFAULT_KEY_RADIX) -> None:
+        self.key = str(key)
+        self.radix = int(radix)
+
+    @classmethod
+    def for_serial(cls, equip_no: str,
+                   radix: int = DEFAULT_KEY_RADIX) -> "EnvelopeCodec":
+        """Codec keyed by a device serial (derives ``generate_secret_key``)."""
+        return cls(generate_secret_key(equip_no), radix=radix)
+
+    # ---------------------------------------------------------------- encode
+    def encode(self, envelope: CommunicationObject) -> bytes:
+        """``toJsonString`` → ``S1 || transform(S2)`` TX bytes."""
+        return encrypt_object(envelope.to_json_string(), self.key,
+                              radix=self.radix)
+
+    # ---------------------------------------------------------------- decode
+    def decode(self, wire: bytes) -> CommunicationObject:
+        """RX bytes → envelope (2nd-comma split, key from ``equipNo``).
+
+        Raises :class:`ValueError` on truncated/non-JSON/undecodable
+        input — callers accumulating a stream treat that as "keep
+        reading".
+        """
+        return CommunicationObject.from_dict(
+            json.loads(decrypt_wire(bytes(wire), radix=self.radix)))
+
+    def try_decode(self, wire: bytes) -> Dict[str, Any] | None:
+        """Like :meth:`decode` but returns the dict, or ``None``."""
+        try:
+            return self.decode(wire).to_dict()
+        except (ValueError, TypeError):
+            return None