From 503b7c8ad80b0c1d064de49fd1a1765be94bcbb5 Mon Sep 17 00:00:00 2001 From: lithium Date: Wed, 15 Feb 2023 15:47:27 +0500 Subject: [PATCH] Init --- README.org | 37 +++++++++++++++++++++ base-config.yaml | 5 +++ boorubot/__init__.py | 1 + boorubot/bot.py | 77 ++++++++++++++++++++++++++++++++++++++++++++ boorubot/db.py | 37 +++++++++++++++++++++ boorubot/image.py | 43 +++++++++++++++++++++++++ build.sh | 3 ++ maubot.yaml | 10 ++++++ 8 files changed, 213 insertions(+) create mode 100644 README.org create mode 100644 base-config.yaml create mode 100644 boorubot/__init__.py create mode 100644 boorubot/bot.py create mode 100644 boorubot/db.py create mode 100644 boorubot/image.py create mode 100755 build.sh create mode 100644 maubot.yaml diff --git a/README.org b/README.org new file mode 100644 index 0000000..0aac519 --- /dev/null +++ b/README.org @@ -0,0 +1,37 @@ +#+STARTUP: overview +#+TITLE: booru-bot +#+LANGUAGE: en +#+OPTIONS: num:nil + +[[https://github.com/maubot/maubot][Maubot]] plugin to fetch random posts from Safebooru. +* Dependencies +** Pip +- maubot +- mautrix +- pymongo +- aiohttp +** MongoDB +* Installaion +- Clone repository +- Run build.sh +- Install mbp file via maubot web interface +- Enter MongoDB credentials in config section +* Usage +** !get +Fetches and sends x amount of pitures that matces specified tags +Examples: +- Fetch 3 random pictures + #+BEGIN_SRC + !g !3 + #+END_SRC +- Fetch 1 picture that match two tags + #+BEGIN_SRC + !g tag_numer_one+tag_number_two + #+END_SRC +- Fetch 5 pictures that match last prompted tags + #+BEGIN_SRC + !g -!5 + #+END_SRC + +** !tags +Toggles sending tags of messages on/off diff --git a/base-config.yaml b/base-config.yaml new file mode 100644 index 0000000..73f16e4 --- /dev/null +++ b/base-config.yaml @@ -0,0 +1,5 @@ +ip: +port: +user: +pass: +db: \ No newline at end of file diff --git a/boorubot/__init__.py b/boorubot/__init__.py new file mode 100644 index 0000000..b3dcdc2 --- /dev/null +++ b/boorubot/__init__.py @@ -0,0 +1 @@ +from .bot import BooruBot diff --git a/boorubot/bot.py b/boorubot/bot.py new file mode 100644 index 0000000..40d53c0 --- /dev/null +++ b/boorubot/bot.py @@ -0,0 +1,77 @@ +from maubot import Plugin, MessageEvent +from maubot.handlers import command +from mautrix.types import ImageInfo +from mautrix.errors import request +from .image import get_image +from .db import store_tags, get_tags, change_tags_listing, get_tags_listing +from mautrix.util.config import BaseProxyConfig, ConfigUpdateHelper +import mimetypes +from typing import Type +import pymongo + + +class Config(BaseProxyConfig): + def do_update(self, helper: ConfigUpdateHelper) -> None: + helper.copy("ip") + helper.copy("port") + helper.copy("user") + helper.copy("pass") + helper.copy("db") + + +class BooruBot(Plugin): + async def start(self) -> None: + self.config.load_and_update() + + @command.new(name="tags", aliases=["t"]) + async def respond2(self, evt: MessageEvent) -> None: + client = pymongo.MongoClient("mongodb://" + str(self.config["user"]) + ":" + str(self.config["pass"]) + "@" + + str(self.config["ip"]) + ":" + str(self.config["port"])) + target = client[str(self.config["db"])]["tags_listing"] + await change_tags_listing(evt.sender, target) + await evt.reply("Tags listing " + await get_tags_listing(evt.sender, target)) + + @command.new(name="get", aliases=["g"]) + @command.argument("prompt", pass_raw=True, required=False) + async def respond(self, evt: MessageEvent, prompt: str) -> None: + arguments = prompt.split('!') + tags = arguments[0] + try: + times = (arguments[1]) + except IndexError: + times = "1" + if times.isdigit(): + times = int(times) + else: + times = 1 + client = pymongo.MongoClient("mongodb://" + str(self.config["user"]) + ":" + str(self.config["pass"]) + "@" + + str(self.config["ip"]) + ":" + str(self.config["port"])) + target = client[str(self.config["db"])]["previous_tags"] + if tags == "-": + tags = await get_tags(evt.sender, target) + if tags == "-": + await evt.reply("No previous requests") + return None + target = client[str(self.config["db"])]["tags_listing"] + tags_listing = await get_tags_listing(evt.sender, target) + for i in range(times): + pic = await get_image(tags) + if pic.valid: + mime_type = mimetypes.guess_type(pic.name)[0] + uri = await self.client.upload_media(pic.image, mime_type=mime_type, filename=pic.name) + try: + await self.client.send_image(evt.room_id, url=uri, file_name=pic.name, + info=ImageInfo(mimetype=mime_type)) + if tags_listing == "enabled": + await evt.respond(pic.tags) + except request.MLimitExceeded: + await evt.reply("Limit exceeded") + else: + await evt.reply("No results") + return None + target = client[str(self.config["db"])]["previous_tags"] + await store_tags(evt.sender, tags, target) + + @classmethod + def get_config_class(cls) -> Type[BaseProxyConfig]: + return Config diff --git a/boorubot/db.py b/boorubot/db.py new file mode 100644 index 0000000..a19b1dd --- /dev/null +++ b/boorubot/db.py @@ -0,0 +1,37 @@ +async def store_tags(sender, tags, target): + query = {"name": str(sender)} + db_entry = target.find_one(query) + if db_entry is None: + target.insert_one({"name": str(sender), "tags": tags}) + else: + target.update_one(query, {"$set": {"tags": tags}}) + + +async def get_tags(sender, target): + query = {"name": str(sender)} + db_entry = target.find_one(query) + if db_entry is None: + return "-" + else: + return db_entry.get('tags') + + +async def change_tags_listing(sender, target): + query = {"name": str(sender)} + db_entry = target.find_one(query) + if db_entry is None: + target.insert_one({"name": str(sender), "tags_listing": "enabled"}) + else: + if db_entry.get('tags_listing') == "enabled": + target.update_one(query, {"$set": {"tags_listing": "disabled"}}) + else: + target.update_one(query, {"$set": {"tags_listing": "enabled"}}) + + +async def get_tags_listing(sender, target): + query = {"name": str(sender)} + db_entry = target.find_one(query) + if db_entry is None: + return "disabled" + else: + return db_entry.get('tags_listing') diff --git a/boorubot/image.py b/boorubot/image.py new file mode 100644 index 0000000..cce000e --- /dev/null +++ b/boorubot/image.py @@ -0,0 +1,43 @@ +import aiohttp +from xml.etree import ElementTree +import random +import math + + +async def make_request(session, url): + async with session.get(url) as response: + return await response.read() + + +class Result: + def __init__(self, image, name, valid): + self.image = image + self.name = name + self.valid = valid + + +async def get_image(tags): + async with aiohttp.ClientSession() as session: + url = ("https://safebooru.org/index.php?page=dapi&s=post&q=index&tags="+tags) + post_front = await make_request(session, url + "&limit=0") + out = Result("", "", True) + tree = ElementTree.fromstring(post_front) + post_count = int(tree.get('count')) + if post_count == 0: + out.valid = False + return out + post_number = random.randrange(post_count) + page_number = math.floor(post_number/100) + url += "&pid=" + str(page_number) + url += "." + if post_number < 10: + url += "0" + url += str(post_number % 100) + page_front = await make_request(session, url) + tree = ElementTree.fromstring(page_front) + image_url = tree.find('post').get('file_url') + image_tags = tree.find('post').get('tags') + out.image = await make_request(session, image_url) + out.name = image_url.replace('https://', '') + out.tags = image_tags + return out diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..d778d07 --- /dev/null +++ b/build.sh @@ -0,0 +1,3 @@ +#!/bin/sh +rm booru_bot.mbp +zip -r booru_bot.mbp boorubot/*.py maubot.yaml base-config.yaml diff --git a/maubot.yaml b/maubot.yaml new file mode 100644 index 0000000..da53a24 --- /dev/null +++ b/maubot.yaml @@ -0,0 +1,10 @@ +maubot: 0.3.1 +id: ru.tsuanagatte.booru_bot +version: 1.0.0 +license: Custom +modules: + - boorubot +main_class: BooruBot +config: true +extra_files: + - base-config.yaml