diff --git a/tgbot/kollagen-bot/main.py b/tgbot/kollagen-bot/main.py index 8bb26b7..ba3826f 100644 --- a/tgbot/kollagen-bot/main.py +++ b/tgbot/kollagen-bot/main.py @@ -2,18 +2,22 @@ from genericpath import isdir import logging import os import subprocess -from typing import Optional +from typing import List, Optional import re + from tempfile import NamedTemporaryFile -from telegram import Update, ForceReply +from telegram import Update from telegram.ext import ( Updater, CommandHandler, DictPersistence, CallbackContext, + defaults, ) +from telegram.parsemode import ParseMode +from parser import SafeArgumentParser, safe_str class KollagenBot: @@ -24,26 +28,78 @@ class KollagenBot: self.kollagen_path = kollagen_path self.base_dir = base_dir - self.updater = Updater(tg_token, persistence=DictPersistence(user_data_json="{}")) + self._init_parser() + + self.updater = Updater( + tg_token, persistence=DictPersistence(user_data_json="{}") + ) dispatcher = self.updater.dispatcher dispatcher.add_handler(CommandHandler("start", self.tg_start)) - dispatcher.add_handler(CommandHandler("list_modes", self.tg_list_modes)) + dispatcher.add_handler(CommandHandler("help", self.tg_help)) dispatcher.add_handler(CommandHandler("generate", self.tg_generate)) dispatcher.add_handler(CommandHandler("g", self.tg_generate)) dispatcher.add_handler(CommandHandler("regenerate", self.tg_regenerate)) dispatcher.add_handler(CommandHandler("r", self.tg_regenerate)) dispatcher.add_error_handler(self.tg_error) + def _init_parser(self): + parser = SafeArgumentParser(prog="/generate", add_help=False) + parser.add_argument( + "directories", + metavar="dir", + type=safe_str, + nargs="*", + default=[self.base_dir] if self.base_dir else [], + help="Directories to process. By default, the entire directory is processed.", + ) + parser.add_argument( + "-m", + dest="mode", + metavar="mode", + choices=self._get_modes(), + nargs="?", + const=True, + help=f"Which collage mode to use. By default, one is chosen at random. When no value is specified, all modes are listed.", + ) + parser.add_argument( + "-n", + dest="num_images", + metavar="N", + type=int, + help=f"How many images to have in a collage. Random (2 None: self.logger.error( @@ -99,7 +180,7 @@ class KollagenBot: def main() -> None: logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - level=logging.INFO, + level=getattr(logging, os.getenv("LOG_LEVEL", "info").upper()), ) tg_token = os.getenv("TG_TOKEN") diff --git a/tgbot/kollagen-bot/parser.py b/tgbot/kollagen-bot/parser.py new file mode 100644 index 0000000..ac5556b --- /dev/null +++ b/tgbot/kollagen-bot/parser.py @@ -0,0 +1,16 @@ +import argparse +import re + + +class ArgumentParserError(RuntimeError): + pass + + +class SafeArgumentParser(argparse.ArgumentParser): + def error(self, message): + raise ArgumentParserError(message) + +def safe_str(val: str): + if re.findall(r'[^\w]', val): + raise RuntimeError("No special characters in arguments allowed!") + return val \ No newline at end of file