Live counter

"""pikoman_bot.py: what a Pikoman bot needs in Python, in one file next to the bot.

- Rest: signed HTTP requests (publish, to one member, stats), each with a token made for that
  one request, as KEYS.md says: the method, the path and a hash of the exact body.
- Bot and Room: a websocket that joins presence channels with join tokens signed by the bot's
  own API key, keeps the member list (and which members are bots), hands client events to your
  handlers, sends, sends to one member and kicks. It reconnects by itself.

Python 3.10+, pip install -r requirements.txt (websockets, BSD-3-Clause; cryptography,
Apache-2.0 OR BSD-3-Clause). Reads PIKOMAN_APP, PIKOMAN_API_KEY and, for tests, PIKOMAN_WS and
PIKOMAN_HTTP from the environment or from a .env file beside the bot.
"""

import asyncio
import base64
import hashlib
import inspect
import json
import os
import random
import signal
import time
import urllib.error
import urllib.parse
import urllib.request

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from websockets.asyncio.client import connect
from websockets.exceptions import ConnectionClosed

LOCAL = ("localhost", "127.0.0.1", "[::1]")


def load_env(path=".env"):
    """KEY=value lines from .env; what the environment already has wins."""
    try:
        with open(path, encoding="utf-8") as file:
            lines = file.read().splitlines()
    except FileNotFoundError:
        return
    for line in lines:
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        name, value = (part.strip() for part in line.split("=", 1))
        value = value.strip("\"'")
        if value and name not in os.environ:
            os.environ[name] = value


load_env()


def env(name, default=""):
    """A setting from the environment; empty counts as not set."""
    return os.environ.get(name) or default


def b64(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def checked(address: str, secure: str) -> str:
    """https:// or wss://; plain http:// or ws:// only to this computer, so keys never go in the clear."""
    parsed = urllib.parse.urlsplit(address)
    if parsed.scheme != secure and not (parsed.scheme == secure[:-1] and parsed.hostname in ("localhost", "127.0.0.1", "::1")):
        raise ValueError(f"{address}: use {secure}:// (plain {secure[:-1]}:// only for localhost)")
    return address.rstrip("/")


class PikomanError(Exception):
    def __init__(self, status, text):
        super().__init__(f"{status}: {text}")
        self.status = status


class Key:
    """An API key's secret, sk_<key id>.<seed>: the id goes in each token's header, the seed signs it."""

    def __init__(self, secret: str):
        if not secret.startswith("sk_") or "." not in secret:
            raise ValueError("PIKOMAN_API_KEY looks like sk_k_….…: the secret the dashboard showed once")
        key_id, seed = secret[3:].split(".", 1)
        self.id = key_id
        self._key = Ed25519PrivateKey.from_private_bytes(base64.urlsafe_b64decode(seed + "=" * (-len(seed) % 4)))

    def sign(self, claims: dict) -> str:
        """An EdDSA JWT with the key id in its header."""
        header = {"alg": "EdDSA", "typ": "JWT", "kid": self.id}
        signed = b64(json.dumps(header).encode()) + "." + b64(json.dumps(claims).encode())
        return signed + "." + b64(self._key.sign(signed.encode()))

    def __repr__(self):
        return f"Key({self.id}, [hidden])"


def target(member) -> dict:
    """A connection id ("c_…"), or {"user_id": "42"} for every connection of that user."""
    if isinstance(member, dict):
        return {"to_user": str(member["user_id"])}
    return {"to": member}


class Rest:
    """The HTTP API, every request signed with the API key and bound to that one request."""

    def __init__(self, app=None, api_key=None, address=None):
        self.app = app or env("PIKOMAN_APP")
        if not self.app:
            raise ValueError("Set PIKOMAN_APP: your app id, app-…")
        self.key = Key(api_key or env("PIKOMAN_API_KEY"))
        self.address = checked(address or env("PIKOMAN_HTTP", "https://next-api.pikoman.se"), "https")

    def request(self, method: str, path: str, data=None) -> dict:
        body = b"" if data is None else json.dumps(data).encode()
        now = int(time.time())
        token = self.key.sign({
            "iss": self.app, "iat": now, "exp": now + 60, "htm": method, "htu": path,
            "bh": b64(hashlib.sha256(body).digest()),
        })
        request = urllib.request.Request(self.address + path, data=body if data is not None else None, method=method, headers={
            "Authorization": "Bearer " + token,
            "Content-Type": "application/json",
        })
        try:
            with urllib.request.urlopen(request, timeout=10) as response:
                text = response.read()
                return json.loads(text) if text else {}
        except urllib.error.HTTPError as e:
            raise PikomanError(e.code, e.read().decode(errors="replace")) from None

    def publish(self, channel: str, event: str, data=None, to=None) -> dict:
        """POST /{channel}/{event}; with `to`, to one member only."""
        path = f"/{urllib.parse.quote(channel, safe='=@,.;')}/{urllib.parse.quote(event, safe='=@,.;')}"
        if to is not None:
            path += "?" + urllib.parse.urlencode(target(to))
        return self.request("POST", path, data)

    def stats(self, channel: str) -> dict:
        return self.request("GET", f"/{urllib.parse.quote(channel, safe='=@,.;')}/stats")


async def call(fn, *args):
    """Runs a handler, plain or async; an error in it is printed and the bot goes on."""
    try:
        result = fn(*args)
        if inspect.isawaitable(result):
            await result
    except Exception as e:  # noqa: BLE001: one handler must never stop the bot
        print(f"handler {getattr(fn, '__name__', fn)}: {e!r}")


def client_name(event: str) -> str:
    return event if event.startswith("client-") else "client-" + event


class Room:
    """A presence channel the bot is in: its members, which of them are bots, and its events."""

    def __init__(self, bot, name: str):
        self.bot = bot
        self.name = name
        self.members = {}  # connection id -> member data, as their join token gave it
        self.bots = {}  # connection id -> the API key id a bot joined with: only the server says
        self._handlers = {}
        self._any = []
        self._joins = []
        self._leaves = []
        self._errors = []

    def on(self, event: str):
        """@room.on("message"): fn(sender, data, info) for that client event (client-message)."""
        def add(fn):
            self._handlers.setdefault(client_name(event), []).append(fn)
            return fn
        return add

    def on_any(self, fn):
        """@room.on_any: fn(sender, event, data, info) for every client event."""
        self._any.append(fn)
        return fn

    def on_join(self, fn):
        """@room.on_join: fn(connection, data) when a member joins."""
        self._joins.append(fn)
        return fn

    def on_leave(self, fn):
        """@room.on_leave: fn(connection, data) when a member leaves."""
        self._leaves.append(fn)
        return fn

    def on_error(self, fn):
        """@room.on_error: fn(code), such as join_refused, not_allowed or rate_limited."""
        self._errors.append(fn)
        return fn

    def is_bot(self, connection: str) -> bool:
        return connection in self.bots

    def member(self, connection: str) -> dict:
        data = self.members.get(connection)
        return data if isinstance(data, dict) else {}

    async def send(self, event: str, data=None):
        """A client event to the other members (at most 10 a second per connection)."""
        await self.bot.send({"op": "send", "channel": self.name, "event": client_name(event), "data": data, "echo": False})

    async def send_to(self, member, event: str, data=None):
        """To one member only: a connection id, or {"user_id": "42"}."""
        await self.bot.send({"op": "send", "channel": self.name, "event": client_name(event), "data": data, "echo": False, **target(member)})

    async def kick(self, member, reason=None, ban_seconds=None):
        """Out of this room: needs the kick right (ban, with ban_seconds)."""
        message = {"op": "kick", "channel": self.name, **target(member)}
        if reason:
            message["reason"] = reason
        if ban_seconds:
            message["ban_seconds"] = ban_seconds
        await self.bot.send(message)

    async def handle(self, message: dict):
        op = message.get("op")
        if op == "members":
            self.members = dict(message.get("members") or {})
            self.bots = dict(message.get("bots") or {})
        elif op == "member_joined":
            self.members[message["id"]] = message.get("data")
            if message.get("bot"):
                self.bots[message["id"]] = message["bot"]
            for fn in self._joins:
                await call(fn, message["id"], message.get("data"))
        elif op == "member_left":
            data = self.members.pop(message["id"], None)
            self.bots.pop(message["id"], None)
            for fn in self._leaves:
                await call(fn, message["id"], data)
        elif op == "client_event":
            sender, event, data = message.get("from"), message.get("event", ""), message.get("data")
            # Set by the server, never by the sender: direct (to this bot alone), user_id, bot.
            info = {k: message[k] for k in ("direct", "user_id", "bot") if k in message}
            for fn in self._handlers.get(event, []):
                await call(fn, sender, data, info)
            for fn in self._any:
                await call(fn, sender, event, data, info)
        elif op == "error":
            print(f"{self.name}: {message.get('code')}")
            for fn in self._errors:
                await call(fn, message.get("code"))


class Bot:
    """One websocket for the bot, in as many rooms as it joins. await bot.run() keeps it going."""

    def __init__(self, name: str, access: int = 3, app=None, api_key=None, address=None):
        self.app = app or env("PIKOMAN_APP")
        if not self.app:
            raise ValueError("Set PIKOMAN_APP: your app id, app-…")
        self.key = Key(api_key or env("PIKOMAN_API_KEY"))
        self.address = checked(address or env("PIKOMAN_WS", "wss://next.pikoman.se"), "wss")
        self.data = {"name": name}  # what the members see next to "bot"
        self.access = access  # 3: may speak (the key needs client_events); 0: listens
        self.rooms = {}
        self.id = None
        self._ws = None
        self._ready = []

    def join(self, channel: str) -> Room:
        room = self.rooms[channel] = Room(self, channel)
        return room

    def on_ready(self, fn):
        """@bot.on_ready: fn() each time the bot is connected and has asked to join its rooms."""
        self._ready.append(fn)
        return fn

    def token(self, channel: str) -> str:
        """A join token for this connection and channel, signed with the bot's API key."""
        return self.key.sign({
            "connection": self.id, "channel": channel, "access": self.access,
            "data": self.data, "exp": int(time.time()) + 300,
        })

    async def send(self, message: dict):
        if self._ws is not None:
            await self._ws.send(json.dumps(message))

    def start(self, *tasks):
        """Runs the bot (and any other coroutines, such as a web server) until Ctrl+C or SIGTERM."""
        def stop(*_):
            raise KeyboardInterrupt

        signal.signal(signal.SIGTERM, stop)  # as from systemd or podman stop: the same as Ctrl+C

        async def main():
            await asyncio.gather(self.run(), *tasks)

        try:
            asyncio.run(main())
        except KeyboardInterrupt:
            print("Stopped.")

    async def run(self):
        delay = 1
        paused = 60  # after 4402: the first wait, doubling to five minutes
        while True:
            try:
                async with connect(f"{self.address}/{self.app}") as ws:
                    self._ws = ws
                    async for raw in ws:
                        message = json.loads(raw)
                        if message.get("op") == "connected":
                            delay = 1
                            self.id = message["id"]
                            for name in self.rooms:
                                await self.send({"op": "join", "channel": name, "token": self.token(name)})
                            for fn in self._ready:
                                await call(fn)
                        elif message.get("channel") in self.rooms:
                            await self.rooms[message["channel"]].handle(message)
                        elif message.get("op") == "error":
                            print("pikoman:", message.get("code"))
            except ConnectionClosed as e:
                code = e.rcvd.code if e.rcvd is not None else None
                if code == 4403:
                    raise SystemExit("The bot's key was deleted.") from None
                if code == 4401:
                    raise SystemExit("The app disconnected the bot.") from None
                if code == 4402:
                    # The key's bot rights are paused (the app's owner is off Pro): not for good.
                    # Try again in minutes, and the bot is back on its own when Pro is.
                    self._ws = None
                    self.id = None
                    print(f"The bot's key is paused until the app is on Pro again; trying again in about {paused // 60} min.")
                    await asyncio.sleep(random.uniform(paused / 2, paused))
                    paused = min(paused * 2, 300)
                    continue
                print("connection closed:", code if code is not None else "no close frame")
            except OSError as e:
                print("could not connect:", e)
            self._ws = None
            self.id = None
            paused = 60
            # Back again after a random wait, longer each time it fails, up to 30 s.
            await asyncio.sleep(random.uniform(0.5, delay))
            delay = min(delay * 2, 30)