matrix/discord-bridge/lotus_presence.py

#!/usr/bin/env python3
"""Discord presence -> Matrix ghosts.

mautrix-discord doesn't bridge presence, so this small service does: it watches
the Discord server over the gateway (GUILD_PRESENCES) and sets presence and a
status message on each member's @discord_<id> ghost, which Lotus already shows.

- Discord bot: the text bridge's own bot, read from its database at startup, so
  there is no extra bot. This is a second gateway session asking only for the
  intents it needs; the bridge's session is untouched.
- Matrix: its own appservice registration (non-exclusive on the ghost namespace;
  the bridge still owns the ghosts). Only ghosts that already exist are updated.
- Synapse drops a user nobody refreshes to offline, so live states are re-sent
  on a timer. That is also the failure mode we want: if this service dies, the
  ghosts go offline rather than showing a stale "online" forever.

Member statuses and activities are never logged.
"""
import asyncio
import logging
import os
import re
import subprocess
import sys
import time
import urllib.parse

import aiohttp

log = logging.getLogger("lotus-presence")

GATEWAY = "wss://gateway.discord.gg/?v=10&encoding=json"
INTENT_GUILDS = 1 << 0
INTENT_GUILD_PRESENCES = 1 << 8
# Close codes where reconnecting cannot help: bad token, disallowed intents.
FATAL_CLOSE = {4004, 4013, 4014}
BUSY = "org.matrix.msc3026.busy"
STATUS_MAP = {"online": "online", "idle": "unavailable", "dnd": BUSY}
MAX_STATUS = 120
# A ghost that doesn't exist yet (never spoke in a bridged room) is retried this often.
UNKNOWN_RETRY_S = 600


def describe_activity(activity):
    """One Discord activity as text, or None. Types per Discord's Activity object."""
    kind = activity.get("type")
    name = (activity.get("name") or "").strip()
    if kind == 4:  # custom status
        emoji = (activity.get("emoji") or {})
        # Custom server emoji have an id and can't be shown as text.
        glyph = emoji.get("name") if not emoji.get("id") else None
        text = (activity.get("state") or "").strip()
        return " ".join(p for p in (glyph, text) if p) or None
    if not name:
        return None
    if kind == 0:
        return f"Playing {name}"
    if kind == 1:
        return f"Streaming {activity.get('details') or name}"
    if kind == 2:
        song, artist = activity.get("details"), activity.get("state")
        if song and artist:
            return f"Listening to {song} by {artist}"
        return f"Listening to {song or name}"
    if kind == 3:
        return f"Watching {name}"
    if kind == 5:
        return f"Competing in {name}"
    return None


def map_presence(status, activities):
    """Discord status + activities -> (Matrix presence, status message)."""
    presence = STATUS_MAP.get(status)
    if presence is None:  # offline / invisible: show nothing about them
        return "offline", ""
    custom = [a for a in activities or [] if a.get("type") == 4]
    others = [a for a in activities or [] if a.get("type") != 4]
    parts = []
    for activity in custom + others[:1]:
        text = describe_activity(activity)
        if text:
            parts.append(text)
    message = " · ".join(parts)
    if len(message) > MAX_STATUS:
        message = message[: MAX_STATUS - 1].rstrip() + "…"
    return presence, message


def ghost_id(discord_id, server_name):
    return f"@discord_{discord_id}:{server_name}"


def bridge_db_env(config_text):
    """libpq settings from the bridge config's database URI (password stays out of argv)."""
    match = re.search(r"^\s*uri:\s*[\"']?(postgres(?:ql)?://[^\"'\s]+)", config_text, re.M)
    if not match:
        raise RuntimeError("no postgres uri in the bridge config")
    uri = urllib.parse.urlparse(match.group(1))
    env = {
        "PGHOST": uri.hostname or "",
        "PGPORT": str(uri.port or 5432),
        "PGUSER": urllib.parse.unquote(uri.username or ""),
        "PGPASSWORD": urllib.parse.unquote(uri.password or ""),
        "PGDATABASE": uri.path.lstrip("/"),
    }
    sslmode = urllib.parse.parse_qs(uri.query).get("sslmode")
    if sslmode:
        env["PGSSLMODE"] = sslmode[0]
    return env


def read_bot_token(config_path):
    """The text bridge's bot token, from its database. Never logged."""
    with open(config_path, encoding="utf-8") as handle:
        env = {**os.environ, **bridge_db_env(handle.read())}
    out = subprocess.run(
        ["psql", "-Atc", "select discord_token from \"user\" where discord_token like 'Bot %'"],
        env=env,
        capture_output=True,
        text=True,
        timeout=30,
        check=False,
    )
    if out.returncode != 0:
        raise RuntimeError("could not read the bot token from the bridge database")
    tokens = [line for line in out.stdout.splitlines() if line]
    if len(tokens) != 1:
        raise RuntimeError(f"expected exactly one bridge bot login, found {len(tokens)}")
    return tokens[0].removeprefix("Bot ")


class GhostPresence:
    """What each ghost should show, pushed to Synapse on change and on a timer."""

    def __init__(self, session, homeserver, as_token, server_name, clock=time.monotonic):
        self.session = session
        self.homeserver = homeserver.rstrip("/")
        self.as_token = as_token
        self.server_name = server_name
        self.clock = clock
        self.live = {}  # discord id -> (presence, message), everything not offline
        self.unknown_until = {}  # discord id -> when to try a missing ghost again

    async def _put(self, discord_id, presence, message):
        user = urllib.parse.quote(ghost_id(discord_id, self.server_name), safe="")
        url = f"{self.homeserver}/_matrix/client/v3/presence/{user}/status?user_id={user}"
        body = {"presence": presence, "status_msg": message}
        headers = {"Authorization": f"Bearer {self.as_token}"}
        try:
            async with self.session.put(
                url, json=body, headers=headers, timeout=aiohttp.ClientTimeout(total=10)
            ) as resp:
                # Synapse answers 403 for a user the appservice hasn't registered.
                if resp.status == 403:
                    self.unknown_until[discord_id] = self.clock() + UNKNOWN_RETRY_S
                elif resp.status >= 400:
                    log.warning("presence update failed: HTTP %s", resp.status)
                else:
                    self.unknown_until.pop(discord_id, None)
        except (aiohttp.ClientError, asyncio.TimeoutError) as err:
            log.warning("presence update failed: %s", type(err).__name__)

    def _skip(self, discord_id):
        return self.unknown_until.get(discord_id, 0) > self.clock()

    async def update(self, discord_id, presence, message):
        known = self.live.get(discord_id)
        if presence == "offline":
            if known is None:
                return
            del self.live[discord_id]
        else:
            if known == (presence, message):
                return
            self.live[discord_id] = (presence, message)
        if not self._skip(discord_id):
            await self._put(discord_id, presence, message)

    async def resync(self, current):
        """A full member snapshot: `current` replaces what we had."""
        for discord_id in set(self.live) - set(current):
            await self.update(discord_id, "offline", "")
        for discord_id, (presence, message) in current.items():
            await self.update(discord_id, presence, message)

    async def refresh(self):
        for discord_id, (presence, message) in list(self.live.items()):
            if not self._skip(discord_id):
                await self._put(discord_id, presence, message)


class FatalGatewayError(Exception):
    pass


class Gateway:
    def __init__(self, session, token, guild_id, ghosts):
        self.session = session
        self.token = token
        self.guild_id = guild_id
        self.ghosts = ghosts
        self.seq = None
        self.session_id = None
        self.resume_url = None
        self.acked = True

    async def _heartbeat(self, ws, interval_s):
        while True:
            await asyncio.sleep(interval_s)
            if not self.acked:
                log.warning("gateway missed a heartbeat ack, reconnecting")
                await ws.close(code=4000)
                return
            self.acked = False
            await ws.send_json({"op": 1, "d": self.seq})

    async def _dispatch(self, kind, data):
        if kind == "READY":
            self.session_id = data["session_id"]
            self.resume_url = data.get("resume_gateway_url")
            log.info("gateway ready")
        elif kind == "GUILD_CREATE" and data.get("id") == self.guild_id:
            current = {}
            for item in data.get("presences") or []:
                presence, message = map_presence(item.get("status"), item.get("activities"))
                if presence != "offline":
                    current[item["user"]["id"]] = (presence, message)
            await self.ghosts.resync(current)
            log.info("synced %d online members", len(current))
        elif kind == "PRESENCE_UPDATE" and data.get("guild_id") == self.guild_id:
            presence, message = map_presence(data.get("status"), data.get("activities"))
            await self.ghosts.update(data["user"]["id"], presence, message)

    async def _run_once(self):
        resuming = bool(self.session_id and self.resume_url)
        url = f"{self.resume_url}/?v=10&encoding=json" if resuming else GATEWAY
        beat = None
        async with self.session.ws_connect(url, heartbeat=None, max_msg_size=0) as ws:
            try:
                async for frame in ws:
                    if frame.type != aiohttp.WSMsgType.TEXT:
                        break
                    msg = frame.json()
                    if msg.get("s") is not None:
                        self.seq = msg["s"]
                    op = msg.get("op")
                    if op == 10:
                        self.acked = True
                        interval = msg["d"]["heartbeat_interval"] / 1000
                        beat = asyncio.create_task(self._heartbeat(ws, interval))
                        if resuming:
                            await ws.send_json({"op": 6, "d": {
                                "token": self.token,
                                "session_id": self.session_id,
                                "seq": self.seq,
                            }})
                        else:
                            await ws.send_json({"op": 2, "d": {
                                "token": self.token,
                                "intents": INTENT_GUILDS | INTENT_GUILD_PRESENCES,
                                "properties": {
                                    "os": "linux",
                                    "browser": "lotus-presence",
                                    "device": "lotus-presence",
                                },
                            }})
                    elif op == 11:
                        self.acked = True
                    elif op == 1:
                        await ws.send_json({"op": 1, "d": self.seq})
                    elif op == 7:
                        return
                    elif op == 9:
                        if not msg.get("d"):
                            self.session_id = None
                        return
                    elif op == 0:
                        await self._dispatch(msg.get("t"), msg.get("d") or {})
            finally:
                if beat:
                    beat.cancel()
            if ws.close_code in FATAL_CLOSE:
                raise FatalGatewayError(f"gateway refused us (close code {ws.close_code})")
            if ws.close_code:
                log.info("gateway closed (code %s)", ws.close_code)

    async def run(self):
        backoff = 1
        while True:
            started = time.monotonic()
            try:
                await self._run_once()
            except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as err:
                log.warning("gateway connection failed: %s", type(err).__name__)
            if time.monotonic() - started > 60:
                backoff = 1
            await asyncio.sleep(backoff)
            backoff = min(backoff * 2, 60)


async def refresh_loop(ghosts, interval_s):
    while True:
        await asyncio.sleep(interval_s)
        await ghosts.refresh()


async def amain():
    token = read_bot_token(os.environ["BRIDGE_CONFIG"])
    async with aiohttp.ClientSession() as session:
        ghosts = GhostPresence(
            session,
            os.environ.get("HOMESERVER_URL", "http://127.0.0.1:8008"),
            os.environ["AS_TOKEN"],
            os.environ["SERVER_NAME"],
        )
        gateway = Gateway(session, token, os.environ["DISCORD_GUILD_ID"], ghosts)
        refresh_s = float(os.environ.get("REFRESH_SECONDS", "20"))
        await asyncio.gather(gateway.run(), refresh_loop(ghosts, refresh_s))


def main():
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    # Discord's own clients log request URLs at DEBUG; keep ours quiet regardless.
    logging.getLogger("aiohttp").setLevel(logging.WARNING)
    try:
        asyncio.run(amain())
    except FatalGatewayError as err:
        log.error("%s", err)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())