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 = 16,
    timeout: float = 5.0,
) -> Dict[str, Any]:
    """Query the aggregate charge percentage over the local TCP socket.

    Tries ``DEVICE_APOWER_INFO`` (0x19) first, then ``BMS`` (0x1D) and
    ``EMS`` (0x1E) 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 in (
            CommunicationCmd.DEVICE_APOWER_INFO,
            CommunicationCmd.BMS,
            CommunicationCmd.EMS,
        ):
            try:
                reply = client.send_rpc(cmd)
            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


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=16, 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")
    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())
    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())