Add --probe2 transport variants (greeting, length prefixes, MQTT, split send)

franklinwh_agate_client/battery.py

@@ -206,6 +206,113 @@ def probe_device(
     return results
 
 
+def _recv_all(sock: socket.socket, idle: float = 4.0) -> bytes:
+    """Drain a socket until idle timeout or close."""
+    out = bytearray()
+    sock.settimeout(idle)
+    while True:
+        try:
+            chunk = sock.recv(65536)
+        except (socket.timeout, TimeoutError):
+            break
+        if not chunk:
+            break
+        out.extend(chunk)
+    return bytes(out)
+
+
+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.envelope import SequenceNumber
+
+    _apower_tx_seq = getattr(_apower_tx, "_seq", None)
+    if _apower_tx_seq is None:
+        _apower_tx_seq = _apower_tx._seq = SequenceNumber()  # type: ignore[attr-defined]
+    env = direct_connect_request(
+        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()
+
+
+#: Minimal MQTT v3.1.1 CONNECT (client id ``probe01``) — rules out / confirms
+#: a local broker on the probed port (CONNACK starts ``20 02``).
+_MQTT_CONNECT = bytes.fromhex("101200044d5154540402003c000670726f62653031")
+
+
+def probe_transport(
+    sn: str,
+    host: str = "10.100.1.1",
+    ports: Tuple[int, ...] = (9000,),
+    key_radix: int = 32,
+    timeout: float = 5.0,
+) -> List[Dict[str, Any]]:
+    """Transport-level probes against one service port.
+
+    The envelope bytes are already verified byte-for-byte against the
+    disassembly; these variants isolate *framing/handshake* behaviour:
+    server-first greeting, immediate raw send with a long read, length
+    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_data
+
+    results: List[Dict[str, Any]] = []
+    for port in ports:
+        tx, envelope = _apower_tx(sn, key_radix)
+
+        def run(label: str, tx_bytes: bytes, pre_read: float = 0.0,
+                post_idle: float = 4.0, split_send: bool = False) -> None:
+            entry: Dict[str, Any] = {
+                "variant": label, "port": port, "tx": tx_bytes,
+                "rx_pre": b"", "rx": b"", "reply": None, "error": None,
+            }
+            try:
+                with socket.create_connection(
+                        (host, port), timeout=timeout) as sock:
+                    if pre_read:
+                        entry["rx_pre"] = _recv_all(sock, idle=pre_read)
+                    if tx_bytes:
+                        if split_send:
+                            half = len(tx_bytes) // 2
+                            sock.sendall(tx_bytes[:half])
+                            sock.settimeout(0.5)
+                            try:
+                                entry["rx_pre"] += sock.recv(65536)
+                            except (socket.timeout, TimeoutError):
+                                pass
+                            sock.sendall(tx_bytes[half:])
+                        else:
+                            sock.sendall(tx_bytes)
+                    entry["rx"] = _recv_all(sock, idle=post_idle)
+            except (ConnectionError, socket.error, TimeoutError) as exc:
+                entry["error"] = f"{type(exc).__name__}: {exc}"
+            if entry["rx"]:
+                decoded = decrypt_data(entry["rx"], radix=key_radix)
+                if decoded:
+                    try:
+                        entry["reply"] = json.loads(decoded)
+                    except ValueError:
+                        entry["reply"] = {"_ undecodable": decoded}
+            results.append(entry)
+
+        run("greet-first/read-then-raw", tx, pre_read=4.0)
+        run("raw/long-read", tx, post_idle=8.0)
+        run("raw/split-send", tx, split_send=True)
+        run("le16-prefix", len(tx).to_bytes(2, "little") + tx)
+        run("be16-prefix", len(tx).to_bytes(2, "big") + tx)
+        run("le32-prefix", len(tx).to_bytes(4, "little") + tx)
+        run("mqtt-connect", _MQTT_CONNECT)
+    return results
+
+
 def main(argv: Optional[List[str]] = None) -> int:
     parser = argparse.ArgumentParser(
         description="Read aGate battery percentage over the device SoftAP socket."
@@ -223,6 +330,9 @@ def main(argv: Optional[List[str]] = None) -> int:
     parser.add_argument("--raw", action="store_true", help="print full reply JSON")
     parser.add_argument("--probe", action="store_true",
                         help="send probe variants and print exact TX/RX bytes")
+    parser.add_argument("--probe2", action="store_true",
+                        help="transport-level probes (greeting, prefixes, "
+                             "MQTT, split send) with exact TX/RX bytes")
     args = parser.parse_args(argv)
 
     if args.show_hotspot:
@@ -254,6 +364,22 @@ def main(argv: Optional[List[str]] = None) -> int:
             if out.get("error"):
                 print(f"  error: {out['error']}")
         return 0
+    if args.probe2:
+        for out in probe_transport(args.sn, host=args.host, ports=ports,
+                                   key_radix=args.key_radix,
+                                   timeout=args.timeout):
+            print(f"--- {out['variant']} -> {args.host}:{out['port']}")
+            print(f"  tx ({len(out['tx'])} B): {out['tx'].hex() or '(none)'}")
+            pre = out.get("rx_pre") or b""
+            if pre:
+                print(f"  rx_pre ({len(pre)} B): {pre.hex()}")
+            rx = out.get("rx") or b""
+            print(f"  rx ({len(rx)} B): {rx.hex() or '(empty)'}")
+            if out.get("reply") is not None:
+                print("  reply: " + json.dumps(out["reply"], default=str))
+            if out.get("error"):
+                print(f"  error: {out['error']}")
+        return 0
     try:
         result = get_battery(args.sn, host=args.host, ports=ports,
                              key_radix=args.key_radix, timeout=args.timeout)