pigway-frame-flow/pigway_frame_flow/integrations/kokonna_official/connection.py

"""Admin account connection; password transient, session persisted as a secret."""

import json
import logging
import os
import tempfile
import threading
from pathlib import Path

from pydantic import Field, SecretStr

from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.integrations.kokonna_official.rpc import OfficialRPC
from pigway_frame_flow.models import StrictModel


class LoginRequest(StrictModel):
    username: str = Field(min_length=1, max_length=120)
    password: SecretStr = Field(min_length=1, max_length=256)


class SelectRequest(StrictModel):
    robot_id: int = Field(gt=0)


class OfficialDevice(StrictModel):
    id: int
    name: str
    model: str | None = None
    model_code: str | None = None
    online: bool = False


class ConnectionState(StrictModel):
    authorized: bool
    username: str
    robot_id: int | None
    mock: bool
    devices: list[OfficialDevice] = Field(default_factory=list)


class Connection:
    def __init__(self, reader, path: Path, secrets: list):
        self.reader, self.path, self.secrets = reader, path, secrets
        self.username = ""
        self.lock = threading.RLock()
        try:
            if path.exists():
                data = json.loads(path.read_text())
                if (
                    not isinstance(data, dict)
                    or not isinstance(data.get("token"), str)
                    or not isinstance(data.get("username"), str)
                    or (
                        data.get("robot_id") is not None
                        and (type(data["robot_id"]) is not int or data["robot_id"] <= 0)
                    )
                    or "robot_id" not in data
                ):
                    raise ValueError()
                reader.rpc = OfficialRPC(data["token"])
                reader._robot_id = data["robot_id"]
                self.username = data["username"]
                self.secrets.append(data["token"])
        except (OSError, ValueError, KeyError, TypeError):
            reader.rpc.disconnect()
            reader._robot_id = None
            logging.getLogger("frame-flow").warning(
                "Official connection configuration unavailable; reconnect required"
            )

    def _save(self, token, username, robot_id):
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(
            mode="w", dir=self.path.parent, delete=False
        ) as f:
            temp = Path(f.name)
            try:
                json.dump(
                    {
                        "token": token,
                        "username": username,
                        "robot_id": robot_id,
                    },
                    f,
                )
                f.flush()
                os.fsync(f.fileno())
                temp.chmod(0o600)
                os.replace(temp, self.path)
            finally:
                temp.unlink(missing_ok=True)

    def _commit(self, rpc, username, robot_id):
        self._save(rpc._token, username, robot_id)
        self.reader.rpc = rpc
        self.reader._robot_id = robot_id
        self.reader._cached = None
        self.username = username
        if rpc._token and rpc._token not in self.secrets:
            self.secrets.append(rpc._token)

    def _query(self, call):
        try:
            return call()
        except FlowError as error:
            if error.code == E.OFFICIAL_AUTH_REQUIRED:
                self.reader.rpc.disconnect()
                self.reader._cached = None
                self.reader._robot_id = None
                try:
                    self._save("", self.username, None)
                except OSError:
                    logging.getLogger("frame-flow").warning(
                        "Expired official connection could not be persisted"
                    )
            raise

    def _devices(self, rpc):
        if self.reader.mock:
            return [
                OfficialDevice(
                    id=1,
                    name="Mock Frame",
                    model="13.3-inch",
                    model_code="B101",
                    online=True,
                )
            ]
        try:
            data = (
                self._query(rpc.devices) if rpc is self.reader.rpc else rpc.devices()
            )["data"]
            return [
                OfficialDevice(
                    id=row["id"],
                    name=row.get("robotNickname") or "Frame",
                    model=(row.get("model") or {}).get("name"),
                    model_code=(row.get("model") or {}).get("code"),
                    online=bool(row.get("online")),
                )
                for row in data
            ]
        except (KeyError, TypeError, ValueError, AttributeError):
            raise FlowError(
                E.CLOUD_UNAVAILABLE, "Official device list unavailable", 502
            ) from None

    def state(self, discovery=False):
        with self.lock:
            return ConnectionState(
                authorized=self.reader.rpc.authorized,
                username=self.username,
                robot_id=self.reader._robot_id,
                mock=self.reader.mock,
                devices=self._devices(self.reader.rpc)
                if discovery and self.reader.rpc.authorized
                else [],
            )

    def login(self, username, password):
        with self.lock:
            rpc = OfficialRPC("mock-session" if self.reader.mock else "")
            if not self.reader.mock:
                rpc.login(username, password)
            devices = self._devices(rpc)
            self._commit(rpc, username, None)
            return ConnectionState(
                authorized=True,
                username=username,
                robot_id=None,
                mock=self.reader.mock,
                devices=devices,
            )

    def select(self, robot_id):
        with self.lock:
            if not self.reader.rpc.authorized:
                raise FlowError(
                    E.CONFIG_REQUIRED, "Official account authorization required", 503
                )
            if not any(d.id == robot_id for d in self._devices(self.reader.rpc)):
                raise FlowError(
                    E.DEVICE_NOT_FOUND, "Official device not in this account", 404
                )
            self._commit(self.reader.rpc, self.username, robot_id)
            return self.state()

    def disconnect(self):
        with self.lock:
            self._commit(OfficialRPC(), "", None)
            return self.state()

    def widgets(self):
        with self.lock:
            return self._query(self.reader.read)