franklinwh-agate-client/franklinwh_agate_client/battery.py

"""Battery charge-state proof of concept.

Direct-connect (SoftAP) read path recovered from the app's Dart AOT layer
(see ``docs/RE_DART_FINDINGS.md``): join the device hotspot, open the local
TCP socket, send the encrypted ``CommunicationObject`` envelope, decrypt the
reply, and pull the aggregate state-of-charge out of it.

No cloud, no MQTT, no Bluetooth needed. Requires L2 reachability to the
aGate gateway (``10.100.1.1``) — i.e. run on the host whose Wi-Fi adapter is
associated to the device AP, or in a container bridged onto ``10.100.1.0/24``.
"""

from __future__ import annotations

import argparse
import json
import logging
import socket
from typing import Any, Dict, List, Optional, Tuple

from .device_local import LocalDeviceClient
from .protocol import CommunicationCmd
from .wifi import derive_hotspot_credentials

log = logging.getLogger(__name__)

#: Ports probed in order. ``9000`` was observed open on one live unit
#: (``docs/RE_DART_FINDINGS.md`` §8); ``18000`` (``0x4650``) is the constant
#: in the app's connect path but was closed on that unit.
SOCKET_PORT_CANDIDATES: Tuple[int, ...] = (9000, 18000)

#: Reply keys searched for the aggregate percentage, in priority order.
SOC_KEYS: Tuple[str, ...] = ("currentBatterySoc", "soc", "batterySoc", "vppSoc")


def _as_dict(data: Any) -> Optional[Dict[str, Any]]:
    if isinstance(data, dict):
        return data
    if isinstance(data, str):
        try:
            parsed = json.loads(data)
        except (ValueError, TypeError):
            return None
        return parsed if isinstance(parsed, dict) else None
    return None


def _walk(obj: Any, key: str) -> List[Any]:
    """Collect every value stored under ``key`` anywhere in nested dicts/lists."""
    found: List[Any] = []
    if isinstance(obj, dict):
        for k, v in obj.items():
            if k == key:
                found.append(v)
            found.extend(_walk(v, key))
    elif isinstance(obj, list):
        for item in obj:
            found.extend(_walk(item, key))
    return found


def _to_float(value: Any) -> Optional[float]:
    try:
        if isinstance(value, bool):
            return None
        num = float(value)
    except (TypeError, ValueError):
        return None
    return num


def extract_soc(reply: Dict[str, Any]) -> Dict[str, Any]:
    """Pull the aggregate charge percentage out of a parsed reply envelope.

    Returns ``{"value": float|None, "source": str|None, "detail": {...}}``.
    ``source`` names the key the value came from; ``detail`` carries the
    supporting fields (``socArray``, ``batteryCount``, charge/discharge
    power) when present.
    """
    data = _as_dict(reply.get("data", reply)) or {}

    for key in SOC_KEYS:
        for container in (data, reply):
            for hit in _walk(container, key):
                num = _to_float(hit)
                if num is not None:
                    return {
                        "value": num,
                        "source": key,
                        "detail": _soc_detail(data),
                    }
    soc_arrays = _walk(data, "socArray")
    for arr in soc_arrays:
        if isinstance(arr, list) and arr:
            nums = [_to_float(v) for v in arr]
            nums = [n for n in nums if n is not None]
            if nums:
                return {
                    "value": sum(nums) / len(nums),
                    "source": "socArray(avg)",
                    "detail": _soc_detail(data),
                }
    return {"value": None, "source": None, "detail": _soc_detail(data)}


def _soc_detail(data: Dict[str, Any]) -> Dict[str, Any]:
    detail: Dict[str, Any] = {}
    for key in (
        "socArray",
        "batteryCount",
        "batteryList",
        "chargePower",
        "dischargePower",
        "currentBatterySoc",
        "soc",
    ):
        hits = _walk(data, key)
        if hits:
            detail[key] = hits[0] if len(hits) == 1 else hits
    # Always record where we looked so failures are debuggable.
    return detail


def get_battery(
    sn: str,
    host: str = "10.100.1.1",
    ports: Tuple[int, ...] = SOCKET_PORT_CANDIDATES,
    key_radix: int = 32,
    timeout: float = 5.0,
) -> Dict[str, Any]:
    """Query the aggregate charge percentage over the local TCP socket.

    Mirrors the app's ``queryAPowerData`` first (``DEVICE_APOWER_INFO`` with
    ``data={"update": 2}``), then ``BMS`` and ``EMS`` as fallbacks, probing
    each candidate port. Returns the ``extract_soc`` dict plus
    ``cmd``/``port``/``reply`` context. Raises the last socket error if no
    port answers.
    """
    last_err: Optional[Exception] = None
    for port in ports:
        client = LocalDeviceClient(
            sn=sn, host=host, socket_port=port,
            key_radix=key_radix, timeout=timeout,
        )
        for cmd, data in (
            (CommunicationCmd.DEVICE_APOWER_INFO, {"update": 2}),
            (CommunicationCmd.BMS, None),
            (CommunicationCmd.EMS, None),
        ):
            try:
                reply = client.send_rpc(cmd, data=data)
            except (ConnectionError, socket.error, TimeoutError) as exc:
                last_err = exc
                log.debug("port %d cmd %s failed: %s", port, cmd, exc)
                break  # port unreachable — try next port
            except Exception as exc:  # envelope-level failure; try next cmd
                last_err = exc
                log.debug("port %d cmd %s error: %s", port, cmd, exc)
                continue
            result = extract_soc(reply)
            result.update({"cmd": str(cmd), "port": port, "reply": reply})
            if result["value"] is not None:
                return result
            last_err = ValueError(f"reply on port {port} cmd {cmd} has no SoC field")
    raise ConnectionError(f"no battery reply via {host}:{list(ports)}") from last_err


#: ``--probe`` request variants: (label, cmd, data).  The first mirrors the
#: app's ``queryAPowerData`` exactly; the rest isolate single envelope
#: decisions (payload present/empty/absent, fallback commands).
PROBE_VARIANTS: Tuple[Tuple[str, CommunicationCmd, Any], ...] = (
    ("apower/update2", CommunicationCmd.DEVICE_APOWER_INFO, {"update": 2}),
    ("apower/none", CommunicationCmd.DEVICE_APOWER_INFO, None),
    ("apower/empty", CommunicationCmd.DEVICE_APOWER_INFO, {}),
    ("queryDeviceInfo/none", CommunicationCmd.QUERY_DEVICE_INFO, None),
    ("bms/none", CommunicationCmd.BMS, None),
    ("ems/none", CommunicationCmd.EMS, None),
)


def probe_device(
    sn: str,
    host: str = "10.100.1.1",
    ports: Tuple[int, ...] = SOCKET_PORT_CANDIDATES,
    key_radix: int = 32,
    timeout: float = 5.0,
) -> List[Dict[str, Any]]:
    """Send every :data:`PROBE_VARIANTS` envelope; return exact TX/RX bytes."""
    results: List[Dict[str, Any]] = []
    for port in ports:
        client = LocalDeviceClient(
            sn=sn, host=host, socket_port=port,
            key_radix=key_radix, timeout=timeout,
        )
        for label, cmd, data in PROBE_VARIANTS:
            try:
                out = client.send_probe(cmd, data=data, timeout=timeout)
            except (ConnectionError, socket.error, TimeoutError) as exc:
                out = {"cmd": str(cmd), "error": f"{type(exc).__name__}: {exc}",
                       "tx": b"", "reply_raw": b"", "reply": None,
                       "envelope": None}
            out.update({"variant": label, "port": port})
            results.append(out)
            if out.get("reply_raw"):
                break  # device answered — deeper variants can wait
    return results


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(
        description="Read aGate battery percentage over the device SoftAP socket."
    )
    parser.add_argument("--sn", required=True, help="20-char device serial")
    parser.add_argument("--host", default="10.100.1.1", help="aGate gateway IP")
    parser.add_argument(
        "--ports", default=",".join(map(str, SOCKET_PORT_CANDIDATES)),
        help="comma-separated socket ports to probe",
    )
    parser.add_argument("--key-radix", type=int, default=32, choices=(16, 32))
    parser.add_argument("--timeout", type=float, default=5.0)
    parser.add_argument("--show-hotspot", action="store_true",
                        help="print derived AP SSID/password and exit")
    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")
    args = parser.parse_args(argv)

    if args.show_hotspot:
        profile = derive_hotspot_credentials(args.sn)
        print(json.dumps(
            {"ssid": profile.ssid, "password": profile.password,
             "gateway_ip": profile.gateway_ip}, indent=2))
        return 0

    ports = tuple(int(p) for p in args.ports.split(",") if p.strip())
    if args.probe:
        from .protocol import generate_secret_key
        print(f"secret key for {args.sn}: {generate_secret_key(args.sn)} "
              f"(radix {args.key_radix})")
        for out in probe_device(args.sn, host=args.host, ports=ports,
                                key_radix=args.key_radix, timeout=args.timeout):
            print(f"--- {out['variant']} -> {args.host}:{out['port']}"
                  f" cmd={out['cmd']}")
            if out.get("envelope") is not None:
                print("  envelope: " + json.dumps(out["envelope"], default=str))
            tx = out.get("tx") or b""
            print(f"  tx ({len(tx)} B): {tx.hex()}")
            rx = out.get("reply_raw") 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))
                soc = extract_soc(out["reply"])
                print(f"  soc: {soc['value']} (source={soc['source']})")
            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)
    except Exception as exc:
        print(f"ERROR: {exc}")
        return 1
    if args.raw:
        print(json.dumps(result["reply"], indent=2, default=str))
    else:
        value = result["value"]
        print(f"{value:g}%" if value is not None else "unknown")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())