~nicoco/jellyfix

53acc102e46b3c2db10e5cb34e5e14126ad8b672 — nicoco 1 year, 2 months ago 19f3388 main
feat: add new example: XMPP bot
1 files changed, 173 insertions(+), 0 deletions(-)

A jellyfix/examples/xmpp_bot.py
A jellyfix/examples/xmpp_bot.py => jellyfix/examples/xmpp_bot.py +173 -0
@@ 0,0 1,173 @@
"""
Notify of new items in an XMPP group chat,

Requires pillow and slixmpp as additional dependencies.
"""

import asyncio
import getpass
import io
import logging
from argparse import ArgumentParser
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path
from typing import Optional

import slixmpp
from PIL import Image

from jellyfin_api_client.api.image import get_item_image
from jellyfin_api_client.models import BaseItemDto, BaseItemKind, ImageType
from jellyfix import JellyfinClient
from jellyfix.examples.keep_original_title import fix


class MyClient(JellyfinClient):
    def __init__(
        self,
        bot_jid: str,
        bot_pass: str,
        muc_jid: str,
        jellyfin_url: str,
        bot_nickname="jellybot",
    ):
        super().__init__(jellyfin_url)

        self.muc_jid = slixmpp.JID(muc_jid)
        self.nickname = bot_nickname
        self.keep_original_titles = True
        self.xmpp = slixmpp.ClientXMPP(bot_jid, bot_pass)
        self.xmpp.register_plugin("xep_0045")
        self.xmpp.register_plugin("xep_0066")
        self.xmpp.register_plugin("xep_0363")
        self.xmpp.add_event_handler("groupchat_message", self.on_message)
        self.xmpp.connect()

        self.upload_file = self.xmpp.plugin["xep_0363"].upload_file
        # let's not block the main thread with image resampling
        self.threads = ThreadPoolExecutor(2)

        self.__notif_queue = list[BaseItemDto]()
        self.__last_add: Optional[datetime] = None
        self.__notify_task: Optional[asyncio.Task] = None

    async def on_message(self, msg: slixmpp.Message):
        if not msg.get_type() != "groupchat":
            return
        # here we could implement commands for the bots

    def send_message(self, body=None, url=None):
        msg = self.xmpp.make_message(mto=self.muc_jid, mtype="groupchat")
        if url:
            msg["oob"]["url"] = url
            msg["body"] = url
        else:
            msg["body"] = body
        msg.send()

    async def start(self):
        await self.xmpp.plugin["xep_0045"].join_muc(self.muc_jid, self.nickname)
        await self.listen(self.handle_event)

    async def __notify_delayed(self):
        await asyncio.sleep(60)
        queue = self.__notif_queue
        series = defaultdict(list)
        for item in queue:
            series[item.series_name] += item
        queue.clear()
        for series, items in series.items():
            n = len(items)
            if n == 1:
                self.send_message(f"1 nouvel épisode de {series}: {items[0].name}")
            else:
                self.send_message(f"{n} nouveaux épisodes de {series}")

    async def handle_event(self, event_type: str, event_data: dict):
        logging.debug("Received event '%s' with data '%s'", event_type, event_data)
        if event_type != "LibraryChanged":
            return
        if not isinstance(event_data, dict):
            return
        added = event_data.get("ItemsAdded", [])
        for item_id in added:
            item = await self.get_item(item_id, sync=False)
            if self.keep_original_titles:
                await fix(self, item)
            await self.notify(item)
        if not self.keep_original_titles:
            return
        updated = event_data.get("ItemsUpdated", [])
        for item_id in updated:
            item = await self.get_item(item_id, sync=False)
            await fix(self, item)

    async def notify(self, item: BaseItemDto):
        if item.type == BaseItemKind.MOVIE:
            body = f"Nouveau film: {item.name} ({item.production_year})"
        elif item.type == BaseItemKind.SERIES:
            body = f"Nouvelle série: {item.name} ({item.production_year})"
        elif item.type == BaseItemKind.EPISODE:
            self.__notif_queue.append(item)
            if self.__notify_task:
                self.__notify_task.cancel()
            self.__notify_task = asyncio.get_running_loop().create_task(self.__notify_delayed())
            return
        else:
            return
        await self.send_thumb(item)
        self.send_message(body)

    async def send_file(self, data, name):
        url = await self.upload_file(filename=Path(f"{name}.jpg"), input_file=data)
        self.send_message(url=url)

    async def send_thumb(self, item: BaseItemDto):
        assert self.client
        assert item.id
        resp = await get_item_image.asyncio_detailed(item.id, ImageType.PRIMARY, client=self.client)
        img = Image.open(io.BytesIO(resp.content))
        await asyncio.get_event_loop().run_in_executor(self.threads, img.thumbnail, (256, 256))
        f = io.BytesIO()
        img.save(f, "JPEG")
        await self.send_file(f, item.name)


def get_parser():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--jellyfin", help="Jellyfin server URL")
    parser.add_argument("--bot-jid", help="JID of your bot")
    parser.add_argument("--bot-pass", help="XMPP password of your bot")
    parser.add_argument("--muc-jid", help="JID of the MUC your bot should join")
    parser.add_argument(
        "--debug",
        action="store_const",
        dest="loglevel",
        const=logging.DEBUG,
        default=logging.INFO,
    )
    parser.add_argument(
        "--keep-original-titles",
        "On item add and update, change their name to the media original title,",
        action="store_true",
    )
    return parser


async def main():
    args = get_parser().parse_args()
    logging.basicConfig(level=args.loglevel)
    client = MyClient(
        args.bot_jid or input("Bot JID?"),
        args.bot_pass or getpass.getpass("Bot XMPP password?"),
        args.muc_jid or input("MUC JID?"),
        args.jellyfin,
    )
    client.keep_original_titles = args.keep_original_titles or False
    await client.start()
    client.threads.shutdown()


asyncio.run(main())