signal-bot/signal_bot/discovery.py

# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""Which groups the bot is in, decided at runtime rather than in .env.

The bot is its own registered Signal account, so the groups it can see are
exactly the groups somebody put it in. Making the operator also write those
ids into .env and restart added a second step to "add the bot" that could
only ever be forgotten. This module keeps the watched set in sync with the
account itself.

The one rule that matters: a *failed* refresh never narrows the set. If the
listing cannot be fetched, or comes back malformed, the previous set stands.
Rebuilding from a partial answer would make the bot conclude it had been
removed from everything -- deaf, and with every group's counts purged.
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, Final

import httpx

from signal_bot.groups import build_alias_map, canonical_group_id
from signal_bot.membership import JoinAttempts, join_pending, verify_allowlist
from signal_bot.roster import RosterCache

if TYPE_CHECKING:
    from signal_bot.config import Config
    from signal_bot.models import IncomingMessage
    from signal_bot.store import Store
    from signal_bot.trust import InviteTrust

logger = logging.getLogger(__name__)

REFRESH_SECONDS: Final = 30.0

Sleeper = Callable[[float], Awaitable[None]]


class GroupDirectory:
    """The live set of groups the bot answers in."""

    def __init__(
        self,
        config: Config,
        client: httpx.AsyncClient,
        store: Store,
        trust: InviteTrust,
        *,
        sleeper: Sleeper | None = None,
    ) -> None:
        """Wire the directory to its config, HTTP client and store.

        Args:
            config: Resolved runtime configuration.
            client: HTTP client for the REST API.
            store: Where per-group state lives, so a group the bot is
                removed from can have its rows dropped.
            trust: Decides whether a discovered group is one the operator
                asked for.
            sleeper: Async sleep used by run(). Injected by the tests.
        """
        self._config = config
        self._client = client
        self._store = store
        self._trust = trust
        self._sleep = sleeper if sleeper is not None else asyncio.sleep
        self._aliases: dict[str, str] = {}
        self._groups: tuple[str, ...] = ()
        self._joins = JoinAttempts()
        # Group ids seen in an envelope that were not in the listing yet. One
        # refresh per *id*, not per unit of time: a global cooldown would
        # make a stranger's message a second earlier swallow the first
        # command in a group that was just added.
        self._refreshed_for: set[str] = set()
        self._members = RosterCache(config.phone_number)

    @property
    def groups(self) -> tuple[str, ...]:
        """Every group currently watched, canonical and in listing order."""
        return self._groups

    def resolve(self, group_id: str | None) -> str | None:
        """Return the canonical id of a watched group, or None.

        This is the single place an envelope's group id becomes a group the
        bot acts on: the scoping gate, the reply address and the per-group
        state key all come from this one answer.

        Args:
            group_id: The id as it arrived, or None for a direct message.

        Returns:
            The canonical id, or None for a direct message or a group the
            bot does not watch.
        """
        if group_id is None:
            return None
        return self._aliases.get(group_id)

    def note_group_update(self, message: IncomingMessage) -> None:
        """Record who changed a group's membership.

        Called before the scoping gate on purpose: the frame that names the
        inviter arrives while the group it describes is still unknown, and
        it is the only authoritative answer to "who added this bot?".

        Args:
            message: Any parsed inbound message.
        """
        if message.group_id is None or not message.is_group_update:
            return
        self._trust.note_update(message.group_id, message.source_number, message.source_uuid)

    async def refresh(self) -> bool:
        """Rebuild the watched set from the account's own group listing.

        Joins anything the bot was merely invited to, then re-reads the
        listing so membership is decided by the API rather than by the join
        endpoint's own optimistic answer.

        Returns:
            True when a listing was actually read. False means the API could
            not be reached and the previous set still stands -- on which no
            caller may conclude the bot is missing from a pinned group.
        """
        entries = await self._fetch()
        if entries is None:
            return False
        if await join_pending(self._client, self._config, entries, self._joins, self._permitted):
            entries = await self._fetch()
            if entries is None:
                return False
        # Before _adopt, deliberately. _adopt returns early once the watched
        # set stops changing -- which is the steady state -- so a roster
        # stashed inside it would freeze at whatever the first sweep saw and
        # never notice anybody joining or leaving again.
        self._members.remember(entries, self._permitted)
        self._adopt(entries)
        return True

    async def refresh_for(self, group_id: str) -> str | None:
        """Refresh once because a message arrived from an unknown group.

        Args:
            group_id: The unrecognised id from the envelope.

        Returns:
            The canonical id if the refresh adopted that group, else None.
        """
        if group_id in self._refreshed_for:
            return None
        self._refreshed_for.add(group_id)
        logger.info("message from unknown group %r; refreshing the listing", group_id)
        await self.refresh()
        return self.resolve(group_id)

    async def verify_allowlist(self) -> None:
        """Refuse to start when GROUP_IDS names a group the bot cannot use."""
        verify_allowlist(self._config, self._groups, self._joins)

    async def run(self, iterations: int | None = None) -> None:
        """Refresh forever, so a group added later needs no restart.

        Args:
            iterations: Stop after this many refreshes. None means forever.
        """
        completed = 0
        while iterations is None or completed < iterations:
            await self._sleep(REFRESH_SECONDS)
            await self.refresh()
            completed += 1

    async def _fetch(self) -> list[dict[str, Any]] | None:
        """Return the group listing, or None when it could not be read.

        None is the "keep what we have" signal. Every failure mode of this
        call -- timeout, 500, a body that is not a list -- must leave the
        watched set alone rather than emptying it.
        """
        try:
            response = await self._client.get(self._config.groups_url)
            response.raise_for_status()
            body = response.json()
        except (httpx.HTTPError, ValueError) as exc:
            logger.warning(
                "group listing unavailable, keeping the %d group(s) already known: %s",
                len(self._groups),
                exc,
            )
            return None
        if not isinstance(body, list):
            logger.warning("group listing was %s, not a list; keeping what we have", type(body))
            return None
        return [entry for entry in body if isinstance(entry, dict)]

    def member_identities(self, group_id: str) -> frozenset[str]:
        """The listing's spelling of every full member of one group."""
        return self._members.identities(group_id)

    def _permitted(self, entry: dict[str, Any]) -> str | None:
        """Return the canonical id of a listing row the bot may watch."""
        raw = entry.get("id") or entry.get("internal_id")
        if not isinstance(raw, str) or not raw:
            return None
        group_id = canonical_group_id(raw)
        if not self._config.allows(group_id):
            return None
        # Two things that are already vouched for. A pinned id was named by
        # the operator in .env; a group already being watched was vouched
        # when it was adopted. Re-testing an adopted group every sweep would
        # let one listing that spells a member as a uuid -- or omits members
        # entirely -- un-watch it and drop its tally, which is the narrowing
        # this module exists to prevent.
        if self._config.is_pinned or group_id in self._groups:
            return group_id
        return group_id if self._trust.allows(group_id, entry) else None

    def _adopt(self, entries: list[dict[str, Any]]) -> None:
        """Replace the watched set from a listing known to be good."""
        groups = tuple(
            group_id
            for entry in entries
            if entry.get("member") and (group_id := self._permitted(entry)) is not None
        )
        if groups == self._groups:
            return
        added = [group_id for group_id in groups if group_id not in self._groups]
        removed = [group_id for group_id in self._groups if group_id not in groups]
        # Built whole and assigned in one statement: the dispatcher reads
        # this map between awaits, and must never see it half-built.
        self._aliases = build_alias_map(groups)
        self._groups = groups
        for group_id in added:
            self._joins.forget(group_id)
            self._refreshed_for.discard(group_id)
            logger.info("now watching %s", group_id)
        for group_id in removed:
            # Dropped, not kept: a re-add should start a fresh tally rather
            # than resurrect counts from whenever the bot was last in there.
            self._store.forget_group(group_id)
            logger.info("no longer a member of %s; its counts were dropped", group_id)