diff --git a/tgbot/config.py b/tgbot/config.py index 002e97c..0fc419f 100644 --- a/tgbot/config.py +++ b/tgbot/config.py @@ -5,3 +5,6 @@ BOORU_API_URL = "http://booru-api:3456/booru" INLINE_QUERY_CACHE_SECONDS = 30 STICKER_SHITHOLE = -1002471390283 + +DANBOORU_USERNAME = "owlrandomshitbot" +DANBOORU_API_KEY = os.getenv("DANBOORU_API_KEY") diff --git a/tgbot/requirements.txt b/tgbot/requirements.txt index 2f74c90..d5c94fa 100644 --- a/tgbot/requirements.txt +++ b/tgbot/requirements.txt @@ -7,4 +7,5 @@ idna==3.10 pillow==11.0.0 python-telegram-bot==21.6 sniffio==1.3.1 -requests==2.32.3 \ No newline at end of file +requests==2.32.3 +Pybooru==4.2.2 \ No newline at end of file diff --git a/tgbot/shit/handlers.py b/tgbot/shit/handlers.py index f71dac3..57fc634 100644 --- a/tgbot/shit/handlers.py +++ b/tgbot/shit/handlers.py @@ -1,9 +1,10 @@ import requests +from pybooru import Danbooru from telegram import Update from telegram.ext import ContextTypes -from tgbot.config import STICKER_SHITHOLE, BOORU_API_URL, INLINE_QUERY_CACHE_SECONDS -from tgbot.shit.hentai import parse_args, create_parser +from tgbot.config import STICKER_SHITHOLE, BOORU_API_URL, INLINE_QUERY_CACHE_SECONDS, DANBOORU_API_KEY +from tgbot.shit.hentai_argparser import parse_args, create_parser from tgbot.shit.render import render_text_on_image import re @@ -108,6 +109,57 @@ async def handle_hentai(update: Update, context: ContextTypes.DEFAULT_TYPE): print(args) + # using danbooru + if args["booru"] in ["db", "dan", "danbooru"]: + + if len(args["tags"]) > 2: + error_response = [ + { + "type": "article", + "id": "1", + "title": "ERROR!", + "description": "None", + "input_message_content": { + "message_text": "None" + } + } + ] + + error_response[0]["description"] = "Danbooru won't let you search more than two tags at once cause " \ + "they're greedy " + error_response[0]["input_message_content"]["message_text"] = error_response[0]["description"] + + return await context.bot.answer_inline_query(update.inline_query.id, error_response, + cache_time=INLINE_QUERY_CACHE_SECONDS) + + client = Danbooru('danbooru', username='owlrandomshitbot', api_key=DANBOORU_API_KEY) + posts = client.post_list(tags=" ".join(args["tags"]), limit=args["limit"], random=args["random"]) + + results = [] + + for idx, post in enumerate(posts): + booru_post_url = "https://danbooru.donmai.us/posts/" + str(post["id"]) + + try: + + results.append( + { + "type": "photo", + "id": "%d" % idx, + "photo_url": post["file_url"], + "thumb_url": post["preview_file_url"], + "caption": "source -> " + booru_post_url + } + ) + + except Exception: + # this post requires a gold account + print(post) + + return await context.bot.answer_inline_query(update.inline_query.id, results, + cache_time=INLINE_QUERY_CACHE_SECONDS) + + # using node booru service response = requests.post(BOORU_API_URL, json=args) if response.status_code == 200: @@ -146,8 +198,8 @@ async def handle_hentai(update: Update, context: ContextTypes.DEFAULT_TYPE): } }) - await context.bot.answer_inline_query(update.inline_query.id, results, - cache_time=INLINE_QUERY_CACHE_SECONDS) + return await context.bot.answer_inline_query(update.inline_query.id, results, + cache_time=INLINE_QUERY_CACHE_SECONDS) else: diff --git a/tgbot/shit/hentai.py b/tgbot/shit/hentai_argparser.py similarity index 100% rename from tgbot/shit/hentai.py rename to tgbot/shit/hentai_argparser.py