diff --git a/tgbot/config.py b/tgbot/config.py index 27965c5..002e97c 100644 --- a/tgbot/config.py +++ b/tgbot/config.py @@ -2,4 +2,6 @@ import os TOKEN = os.getenv("TOKEN") BOORU_API_URL = "http://booru-api:3456/booru" +INLINE_QUERY_CACHE_SECONDS = 30 + STICKER_SHITHOLE = -1002471390283 diff --git a/tgbot/shit/handlers.py b/tgbot/shit/handlers.py index 77d5f4f..f71dac3 100644 --- a/tgbot/shit/handlers.py +++ b/tgbot/shit/handlers.py @@ -1,10 +1,19 @@ import requests -from telegram import Update, InlineQueryResultPhoto +from telegram import Update from telegram.ext import ContextTypes -from tgbot.config import STICKER_SHITHOLE, BOORU_API_URL -from tgbot.shit.hentai import QueryParams +from tgbot.config import STICKER_SHITHOLE, BOORU_API_URL, INLINE_QUERY_CACHE_SECONDS +from tgbot.shit.hentai import parse_args, create_parser from tgbot.shit.render import render_text_on_image +import re + + +def escape_markdown(text): + """ + Escapes special characters for MarkdownV2 formatting. + """ + escape_chars = r"[_*[\]()~`>#+-=|{}.!]" + return re.sub(f"([{escape_chars}])", r"\\\1", text) async def handle_xitter(update: Update, context: ContextTypes.DEFAULT_TYPE): @@ -44,30 +53,62 @@ async def handle_red_ebalo(update: Update, context: ContextTypes.DEFAULT_TYPE): async def handle_hentai(update: Update, context: ContextTypes.DEFAULT_TYPE): + query = update.inline_query.query.replace("hentai", "").replace("—", "--") - error_response = [ - { - "type": "article", - "id": "1", - "title": "ERROR OCCURRED! That means, at least one of us is a retard.", - "description": "None", - "input_message_content": { - "message_text": "None" - } - } - ] - - query = update.inline_query.query + parser = create_parser() try: - params = QueryParams.from_query_str(query).__dict__ - except ValueError as e: - print(e) - error_response[0]["description"] = "error parsing query" - error_response[0]["input_message_content"]["message_text"] = "error parsing query" - return await context.bot.answer_inline_query(update.inline_query.id, error_response) - response = requests.post(BOORU_API_URL, json=params) + args = parse_args(query, parser) + + except RuntimeError as e: + + if getattr(parser, 'help_message', None): + help_message = parser.help_message + + escaped_help_message = escape_markdown(help_message) + + help_response = [ + { + "type": "article", + "id": "1", + "title": "Help message", + "description": "Print help", + "input_message_content": { + "message_text": f"```\n{escaped_help_message}\n```", + "parse_mode": "MarkdownV2", + } + } + ] + + return await context.bot.answer_inline_query(update.inline_query.id, help_response, + cache_time=INLINE_QUERY_CACHE_SECONDS) + + else: + + error_message = str(e) + + escaped_error_message = escape_markdown(error_message) + + error_response = [ + { + "type": "article", + "id": "1", + "title": "Invalid command!", + "description": error_message, + "input_message_content": { + "message_text": f"```\n{escaped_error_message}\n```", + "parse_mode": "MarkdownV2", + } + } + ] + + return await context.bot.answer_inline_query(update.inline_query.id, error_response, + cache_time=INLINE_QUERY_CACHE_SECONDS) + + print(args) + + response = requests.post(BOORU_API_URL, json=args) if response.status_code == 200: @@ -79,24 +120,51 @@ async def handle_hentai(update: Update, context: ContextTypes.DEFAULT_TYPE): booru_post_url = post["booru"]["domain"] + \ post["booru"]["site"]["api"]["postView"] + post["id"] - results.append( - InlineQueryResultPhoto( - id="%d" % idx, - photo_url=post["fileUrl"], - thumbnail_url=post["previewUrl"], - title="IMAGE #%d" % idx, - description="Image from booru", - caption="source -> " + booru_post_url - ) - ) + if not args["spoiler"]: - await context.bot.answer_inline_query(update.inline_query.id, results) + results.append( + { + "type": "photo", + "id": "%d" % idx, + "photo_url": post["fileUrl"], + "thumb_url": post["previewUrl"], + "caption": "source -> " + booru_post_url + } + ) + + else: + + results.append({ + "type": "photo", + "id": "%d" % idx, + "photo_url": post["fileUrl"], + "thumb_url": post["previewUrl"], + "caption": "source -> " + booru_post_url, + "input_message_content": { + "message_text": "image with tags = " + str(args["tags"]) + " -> " + booru_post_url, + "has_spoiler": True + } + }) + + await context.bot.answer_inline_query(update.inline_query.id, results, + cache_time=INLINE_QUERY_CACHE_SECONDS) else: - print(response.text) + error_response = [ + { + "type": "article", + "id": "1", + "title": "ERROR!", + "description": "None", + "input_message_content": { + "message_text": "None" + } + } + ] error_response[0]["description"] = "booru error" error_response[0]["input_message_content"]["message_text"] = response.text - return await context.bot.answer_inline_query(update.inline_query.id, error_response) + return await context.bot.answer_inline_query(update.inline_query.id, error_response, + cache_time=INLINE_QUERY_CACHE_SECONDS) diff --git a/tgbot/shit/hentai.py b/tgbot/shit/hentai.py index b68a105..cbd73aa 100644 --- a/tgbot/shit/hentai.py +++ b/tgbot/shit/hentai.py @@ -1,46 +1,76 @@ -import re +import argparse +import io +import shlex -def combine_whitespaces(s: str) -> str: - return re.compile(r"\s+").sub(" ", s).strip() +class CustomArgParser(argparse.ArgumentParser): + def __init__( + self, + *args, **kwargs, + ): + super().__init__(*args, **kwargs) + self.help_message = None + + def print_help(self, file=None): + # Store the help message in a buffer instead of printing + if file is None: + help_buffer = io.StringIO() + super().print_help(file=help_buffer) + self.help_message = help_buffer.getvalue() + help_buffer.close() + else: + super().print_help(file=file) + + def parse_args(self, args=None, namespace=None): + # Check for --help manually to avoid stdout output + if '--help' in args: + self.print_help() + raise RuntimeError("Help requested") + return super().parse_args(args, namespace) + + def exit(self, status=0, message=None): + raise RuntimeError(message) + + def error(self, message): + raise RuntimeError(f"Error: {message}") -class QueryParams: - def __init__(self, booru: str, tags: list[str], random: bool = True, limit: int = 10): - self.booru = booru - self.tags = tags - self.random = random - self.limit = limit +def create_parser(): + parser = CustomArgParser( + description="A command to search random stuff on various image boorus. Feel free to use it irresponsibly!", + usage="@owlrandomshitbot hentai BOORU_NAME --tags TAG1,TAG2,TAG3 -l LIMIT --random -s", + ) - @classmethod - def from_query_str(cls, query: str): - query = query.replace("hentai ", "") + parser.add_argument("booru", type=str, + choices=['e6', 'e621', 'e9', 'e926', 'hh', 'hypno', 'hypnohub', 'db', 'dan', 'danbooru', 'kc', + 'konac', 'kcom', 'kn', 'konan', 'knet', 'yd', 'yand', 'yandere', 'gb', 'gel', + 'gelbooru', + 'r34', 'rule34', 'sb', 'safe', 'safebooru', 'tb', 'tbib', 'big', 'xb', 'xbooru', 'pa', + 'paheal', 'dp', 'derp', 'derpi', 'derpibooru', 'rb', 'realbooru'], + help="Booru to search. For context, this command uses this library" + " -> https://github.com/AtoraSuunva/booru" + ) - query = combine_whitespaces(query) + parser.add_argument("--tags", "-t", + dest="tags", type=str, help="Tags, comma separated. If query contains tags with spaces, " + "e.g. 'downvote bait', put it in quotes.") - list_ = query.split(" ") + parser.add_argument("--random", "-r", + dest="random", action='store_true', default=False, + help="Randomize the output, default=False") - booru = list_[0] + parser.add_argument("--limit", "-l", + dest="limit", type=int, default=20, help="Limit the number of output posts, default=20") - list_.remove(booru) + parser.add_argument("--spoiler", "-s", dest="spoiler", + action="store_true", help="Send the image with a spoiler.") - tags = [] - random = False - limit = 10 + return parser - for e in list_: - if "random=" in e: # this is 'random' argument - random = e.replace("random=", "").lower() == "true" - continue - if "limit=" in e: # this is 'limit' argument - limit = int(e.replace("limit=", "")) - continue +def parse_args(args: str, parser: CustomArgParser): + args = parser.parse_args(shlex.split(args)).__dict__ - tags.append(e) + args["tags"] = args["tags"].split(",") - return cls(booru, tags, random, limit) - - def __str__(self): - return "QueryParams(booru=" + self.booru + ", tags=" + str(self.tags) + ", random=" + str(self.random) + \ - ", limit=" + str(self.limit) + ")" + return args