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"]
replies_enabled = secrets["groupme"].get("replies_enabled", True)
groupme_admins = secrets["admins"]["origin"]
group_admins = secrets["admins"]["channel"]
imgur_fallback = secrets["groupme"].get("imgur_fallback", False)
imgur_token = secrets["imgur"]["token"] if imgur_fallback else None

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": str(int(att_body["reply_id"]) - 1),
                        },
                    ) 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 on_message(self, incoming: GroupMeMessage):
        parsed = await self.parse(incoming)
        if parsed.sender_name.lower() == "rollbot":
            return
        await self.respond(rollbot.Response.from_message(
            parsed,
            f"Hey {parsed.sender_name}, GroupMe is a dying chat platform, please consider using Discord instead!",
            [rollbot.Attachment("reply", parsed.message_id)]
        ))

    async def upload_image(self, data: bytes) -> tuple[str, str]:
        try:
            return await self.upload_image_groupme(data), ""
        except:
            self.context.logger.exception("GroupMe image upload failure")
            if not imgur_fallback:
                self.context.logger.info("Imgur fallback disabled")
                raise
            self.context.logger.info("Attempting imgur fallback")
            return await self.upload_image_imgur(data), "GroupMe rejected image upload, falling back to Imgur...\n"

    async def upload_image_groupme(self, data: bytes) -> str:
        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 upload_image_imgur(self, data: bytes) -> str:
        async with self.context.request.post(
            "https://api.imgur.com/3/image",
            headers={
                "Authorization": f"Client-ID {imgur_token}",
            },
            data={"image": data},
        ) as upload:
            upload.raise_for_status()
            return (await upload.json())["data"]["link"]

    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):
                            url, note = await self.upload_image(att.body)
                            attachments.append(
                                {
                                    "type": "image",
                                    "url": url,
                                }
                            )
                            message += note
                        else:
                            attachments.append({"type": "image", "url": att.body})
                    if replies_enabled and 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)