home_control/lib/_zyxel.py

"""Small client for the encrypted Zyxel OPAL/DAL web API."""

from __future__ import annotations

import argparse
import base64
import json
import logging
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit

import requests
import yaml
from Crypto.Cipher import AES, PKCS1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Util.Padding import pad, unpad


LOG = logging.getLogger("home_control.zyxel")


class ZyxelError(RuntimeError):
    """The router rejected a request or returned an unexpected response."""


def normalize_mac(value: str) -> str:
    compact = "".join(character for character in value if character.isalnum()).lower()
    if len(compact) != 12 or any(character not in "0123456789abcdef" for character in compact):
        raise ValueError(f"invalid MAC address: {value!r}")
    return ":".join(compact[index : index + 2] for index in range(0, 12, 2))


@dataclass(frozen=True)
class LanHost:
    mac: str
    active: bool
    hostname: str = ""
    ip_address: str = ""


class ZyxelRouter:
    """Read-only client for Zyxel routers using the OPAL DAL API."""

    def __init__(
        self,
        host: str,
        username: str,
        password: str,
        *,
        timeout: float = 10,
        verify_tls: bool = True,
        hosts_oid: str = "lanhosts",
        debug: bool = False,
    ) -> None:
        if "://" not in host:
            host = f"http://{host}"
        parsed = urlsplit(host)
        if parsed.scheme not in {"http", "https"} or not parsed.netloc:
            raise ValueError("zyxel.host must be a hostname/IP, or an http(s) URL")
        self.url = host.rstrip("/")
        self.username = username
        self.password = password
        self.timeout = timeout
        self.verify_tls = verify_tls
        self.hosts_oid = hosts_oid
        self.debug = debug
        self.session = requests.Session()
        self._aes_key: bytes | None = None
        self._session_key: str | None = None
        self._has_logged_connection = False

    def __enter__(self) -> "ZyxelRouter":
        self.login()
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    def _request(self, method: str, path: str, **kwargs: Any) -> requests.Response:
        kwargs.setdefault("timeout", self.timeout)
        kwargs.setdefault("verify", self.verify_tls)
        try:
            response = self.session.request(method, f"{self.url}{path}", **kwargs)
            response.raise_for_status()
            if self.debug:
                LOG.info("Zyxel HTTP call: %s %s -> %s", method, path, response.status_code)
            return response
        except requests.RequestException as error:
            raise ZyxelError(f"Zyxel request failed: {error}") from error

    @staticmethod
    def _json(response: requests.Response) -> dict[str, Any]:
        try:
            value = response.json()
        except (requests.JSONDecodeError, ValueError) as error:
            raise ZyxelError("Zyxel returned a non-JSON response") from error
        if not isinstance(value, dict):
            raise ZyxelError("Zyxel returned an unexpected JSON response")
        return value

    def login(self) -> None:
        self._request("GET", "/GetInfoNoLogin")
        public_key = self._json(self._request("GET", "/getRSAPublickKey")).get("RSAPublicKey")
        if not public_key:
            raise ZyxelError("Zyxel did not return an RSA public key")

        self._aes_key = os.urandom(32)
        iv = os.urandom(32)
        login = {
            "Input_Account": self.username,
            "Input_Passwd": base64.b64encode(self.password.encode()).decode(),
            "currLang": "en",
            "RememberPassword": 0,
        }
        cipher = AES.new(self._aes_key, AES.MODE_CBC, iv[:16])
        content = cipher.encrypt(pad(json.dumps(login, separators=(",", ":")).encode(), 16))
        rsa = PKCS1_v1_5.new(RSA.import_key(public_key.encode()))
        encrypted_key = rsa.encrypt(base64.b64encode(self._aes_key))
        payload = {
            "content": base64.b64encode(content).decode(),
            "key": base64.b64encode(encrypted_key).decode(),
            "iv": base64.b64encode(iv).decode(),
        }
        result = self._decrypt(self._json(self._request("POST", "/UserLogin", json=payload)))
        if result.get("result") != "ZCFG_SUCCESS" or not result.get("sessionkey"):
            raise ZyxelError(f"Zyxel login failed: {result.get('result', 'unknown error')}")
        self._session_key = str(result["sessionkey"])
        self._log_first_connection()

    def _log_first_connection(self) -> None:
        if not self._has_logged_connection:
            LOG.info("Connected successfully to Zyxel router at %s", self.url)
            self._has_logged_connection = True

    def _decrypt(self, envelope: dict[str, Any]) -> dict[str, Any]:
        if "content" not in envelope or "iv" not in envelope:
            raise ZyxelError(str(envelope.get("result") or "unencrypted Zyxel response"))
        if self._aes_key is None:
            raise ZyxelError("not logged in")
        try:
            iv = base64.b64decode(envelope["iv"])[:16]
            encrypted = base64.b64decode(envelope["content"])
            clear = unpad(AES.new(self._aes_key, AES.MODE_CBC, iv).decrypt(encrypted), 16)
            result = json.loads(clear)
        except (ValueError, KeyError, json.JSONDecodeError) as error:
            raise ZyxelError("could not decrypt Zyxel response") from error
        if not isinstance(result, dict):
            raise ZyxelError("Zyxel returned an unexpected encrypted response")
        return result

    def dal_get(self, oid: str) -> dict[str, Any]:
        if not self._session_key:
            self.login()
        response = self._request(
            "GET", "/cgi-bin/DAL", params={"oid": oid, "sessionkey": self._session_key}
        )
        result = self._decrypt(self._json(response))
        if self.debug:
            LOG.info("Zyxel DAL call: oid=%s -> %s", oid, result.get("result", "no result"))
        return result

    def get_lan_hosts(self) -> list[LanHost]:
        result = self.dal_get(self.hosts_oid)
        if result.get("result") not in (None, "ZCFG_SUCCESS"):
            raise ZyxelError(f"lanhosts query failed: {result.get('result')}")

        def host_records(value: Any):
            """Yield host dictionaries from flat and nested DAL response shapes."""
            if isinstance(value, dict):
                if any(
                    key in value
                    for key in ("PhysAddress", "physAddress", "MACAddr", "MacAddress")
                ):
                    yield value
                    return
                for child in value.values():
                    yield from host_records(child)
            elif isinstance(value, list):
                for child in value:
                    yield from host_records(child)

        hosts: list[LanHost] = []
        for item in host_records(result.get("Object", [])):
            raw_mac = next((item.get(key) for key in ("PhysAddress", "physAddress", "MACAddr", "MacAddress") if item.get(key)), None)
            if not raw_mac:
                continue
            try:
                mac = normalize_mac(str(raw_mac))
            except ValueError:
                continue
            raw_active = next((item.get(key) for key in ("Active", "active", "Enable", "enable") if key in item), True)
            active = raw_active if isinstance(raw_active, bool) else str(raw_active).lower() in {"1", "true", "yes", "active", "enabled"}
            hosts.append(LanHost(
                mac=mac,
                active=active,
                hostname=str(item.get("HostName") or item.get("hostName") or ""),
                ip_address=str(item.get("IPAddress") or item.get("ipAddress") or item.get("IPAddr") or ""),
            ))
        return hosts

    def connection_states(self, macs: list[str]) -> dict[str, bool]:
        """Check several MAC addresses using one router query."""
        wanted = [normalize_mac(mac) for mac in macs]
        active_macs = {host.mac for host in self.get_lan_hosts() if host.active}
        states = {mac: mac in active_macs for mac in wanted}
        if self.debug:
            for mac, connected in states.items():
                LOG.info(
                    "Zyxel presence result: mac=%s -> %s",
                    mac,
                    "connected" if connected else "absent",
                )
        return states

    def is_connected(self, mac: str) -> bool:
        wanted = normalize_mac(mac)
        return self.connection_states([wanted])[wanted]

    def close(self) -> None:
        if self._session_key:
            try:
                self._request("GET", "/cgi-bin/UserLogout", params={"sessionkey": self._session_key})
            except ZyxelError:
                pass
        self._session_key = None
        self._aes_key = None
        self.session.close()


def _load_cli_config(path: Path) -> tuple[dict[str, Any], str]:
    try:
        with path.open(encoding="utf-8") as stream:
            root = yaml.safe_load(stream)
        arrival = root["arrival_detection"]
        return arrival["zyxel"], str(arrival["devices"][0]["mac"])
    except (OSError, TypeError, KeyError, yaml.YAMLError) as error:
        raise ZyxelError(f"could not load configuration from {path}: {error}") from error


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Check whether a MAC address is connected to the configured Zyxel router."
    )
    parser.add_argument(
        "--config",
        type=Path,
        default=Path(__file__).resolve().parents[1] / "service.yaml",
        help="service configuration file (default: project service.yaml)",
    )
    parser.add_argument(
        "--mac",
        help="MAC address to check (default: first arrival_detection.devices entry)",
    )
    args = parser.parse_args()

    try:
        config, configured_mac = _load_cli_config(args.config)
        mac = normalize_mac(args.mac or configured_mac)
        router = ZyxelRouter(
            config["host"],
            config["username"],
            config["password"],
            timeout=float(config.get("timeout", 10)),
            verify_tls=bool(config.get("verify_tls", True)),
            hosts_oid=str(config.get("hosts_oid", "lanhosts")),
            debug=bool(config.get("debug", False)),
        )
        try:
            connected = router.is_connected(mac)
        finally:
            router.close()
        print(f"{mac} {'connected' if connected else 'absent'}")
        return 0
    except (ValueError, ZyxelError) as error:
        print(f"error: {error}", file=sys.stderr)
        return 2


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