diff --git a/franklinwh_agate_client/battery.py b/franklinwh_agate_client/battery.py index a2295bc..6598723 100644 --- a/franklinwh_agate_client/battery.py +++ b/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)