rollbot4/drivers/groupme_driver.py

from __future__ import annotations

from datetime import datetime
import json
import traceback
import asyncio
import logging
import os
import time

import toml
from fastapi import FastAPI, BackgroundTasks, Response
from pydantic import BaseModel

import rollbot
from commands import config

logging.config.fileConfig("logging.conf", disable_existing_loggers=False)

with open(os.environ.get("SECRET_FILE", "secrets.toml"), "r") as sfile:
    secrets = toml.load(sfile)
database_file = os.environ.get("DATABASE_FILE", secrets["database_file"])
groupme_bots = secrets["groupme"]["bots"]
groupme_token = secrets["groupme"]["token"]
groupme_admins = secrets["admins"]["origin"]
group_admins = secrets["admins"]["channel"]

max_msg_len = 1000
split_text = "\n..."
msg_slice = max_msg_len - len(split_text)


class GroupMeMessage(BaseModel):
    id: str
    sender_id: str
    group_id: str
    name: str
    text: str
    created_at: int
    attachments: list[dict[str, str]]


class GroupMeBot(rollbot.Rollbot[GroupMeMessage]):
    def __init__(self):
        super().__init__(config.extend(rollbot.CommandConfiguration(bangs=("!",))), database_file)

    def read_config(self, key: str) -> Any:
        cfg = secrets
        for part in key.split("."):
            cfg = cfg.get(part, None)
            if cfg is None:
                return None
        return cfg

    def _convert_attachment(self, group_id, att_type, att_body):
        if att_type == "reply":
            async def fetch_reply():
                try:
                    async with self.context.request.get(
                        f"https://api.groupme.com/v3/groups/{group_id}/messages",
                        headers={
                            "Content-Type": "application/json",
                        },
                        params={
                            "token": groupme_token,
                            "limit": 1,
                            "after_id": att_body["base_reply_id"],
                        },
                    ) as resp:
                        msg = (await resp.json())["response"]["messages"][0]
                    return json.dumps(msg)
                except:
                    self.context.logger.exception("Failed to look up attached message")
            return rollbot.Attachment(att_type, fetch_reply)
        return rollbot.Attachment(att_type, json.dumps(att_body))

    async def parse(self, msg: GroupMeMessage):
        return rollbot.Message(
            origin_id="GROUPME",
            channel_id=msg.group_id,
            sender_id=msg.sender_id,
            timestamp=datetime.fromtimestamp(msg.created_at),
            origin_admin=msg.sender_id in groupme_admins,
            channel_admin=msg.sender_id in group_admins.get(msg.group_id, ()),
            sender_name=msg.name,
            text=msg.text,
            attachments=[self._convert_attachment(msg.group_id, d["type"], d) for d in msg.attachments],
            message_id=msg.id,
        )

    async def upload_image(self, data: bytes):
        async with self.context.request.post(
            "https://image.groupme.com/pictures",
            headers={
                "Content-Type": "image/png",
                "X-Access-Token": groupme_token,
            },
            data=data,
        ) as upload:
            upload.raise_for_status()
            return (await upload.json())["payload"]["url"]

    async def post_message(self, bot_id: str, text: str, attachments: list[dict[str, str]]):
        body = {
            "bot_id": bot_id,
            "text": text,
            "attachments": attachments,
        }
        self.context.logger.info(f"Sending: {body}")
        result = await self.context.request.post(
            "https://api.groupme.com/v3/bots/post",
            json=body,
            timeout=10,
        )
        self.context.logger.info(f"Received: {result.status} - {await result.text()}")

    async def respond(self, res: rollbot.Response):
        if res.cause is not None and (proc_time := time.time() - res.cause.received_at) < 1:
            # sleep for a moment to make groupme not misorder messages
            await asyncio.sleep(1 - proc_time)

        if res.origin_id != "GROUPME":
            self.context.logger.error(f"Unable to respond to {res.origin_id}")
            return

        bot_id = groupme_bots.get(res.channel_id, None)
        if bot_id is None:
            self.context.logger.error(f"Unable to respond to group {res.channel_id} in GroupMe")
            return

        message = ""
        attachments = []
        try:
            if res.attachments is not None:
                for att in res.attachments:
                    if att.name == "image":
                        if isinstance(att.body, bytes):
                            attachments.append(
                                {
                                    "type": "image",
                                    "url": await self.upload_image(att.body),
                                }
                            )
                        else:
                            attachments.append({"type": "image", "url": att.body})
                    if att.name == "reply":
                        if att.body is None or not isinstance(att.body, str):
                            raise ValueError("Invalid reply body type, must be message ID")
                        attachments.append({
                            "type": "reply",
                            "base_reply_id": att.body,
                            "reply_id": att.body,
                        })
                         
        except:
            self.context.debugging = "".join(traceback.format_exc())
            message += "Failed to upload one or more attachments!\n"
            self.context.logger.exception("Failed to upload attachment")

        if res.text is not None:
            message += res.text

        msgs = []
        while len(message) > max_msg_len:
            msgs.append(message[:msg_slice] + split_text)
            message = message[msg_slice:]
        msgs.append(message)

        await self.post_message(bot_id, msgs[0], attachments)
        for m in msgs[1:]:
            await asyncio.sleep(0.1)
            await self.post_message(bot_id, m, [])


app = FastAPI()
bot = GroupMeBot()


@app.on_event("startup")
async def startup():
    await bot.on_startup()


@app.on_event("shutdown")
async def shutdown():
    await bot.on_shutdown()


@app.post("/", status_code=204)
async def receive(message: GroupMeMessage, tasks: BackgroundTasks):
    tasks.add_task(bot.on_message, message)
    return Response(status_code=204)


@app.get("/health", status_code=204)
def health():
    return Response(status_code=204)