import re import logging import peewee import datetime from db.queries import skip_day_put, skip_day_delete, skip_days_upcoming, user_add, user_get, domsubusers_add, domsubusers_delete from db.models import OrdersPool from .telegram import Telegram, TelegramCommand from settings import FLASK_URL logger = logging.getLogger(__name__) CALLBACK_PREFIX_SKIP_DAY_ADD = "skip_day_add_" class StartCommand(TelegramCommand): command_text = '/start' description = "Start a session with the bot" async def exec_inner(self, text, update, *args, **kwargs): username = update['message']['chat']['username'] chat_id = update['message']['chat']['id'] user_add(username, chat_id) yield f"Welcome to Gear Orders bot.\n\nEdit your orders at {FLASK_URL}\n\nAdd a dom with `/dom_add ` or ask a sub to add you." class DomAddCommand(TelegramCommand): command_text = "/dom_add" command_regex = re.compile(r"^\/dom_add( (?P@?\w+))$") description = "Add a dom to your profile" async def exec_inner(self, text, update, session, *args, **kwargs): dom_username = text.split(' ')[1] if dom_username.startswith('@'): dom_username = dom_username[1:] sub_username = update['message']['chat']['username'] sub = user_get(sub_username) try: dom = user_get(dom_username) except: yield "I'm sorry that user isn't registered with me." return try: domsubusers_add(sub, dom) except peewee.IntegrityError: yield f"@{dom_username} is already on your list of doms" return yield f"Successfully added {dom_username} to your list of doms" t = Telegram(session) await t.message_send(dom.telegram_chat_id, f"@{sub_username} has added you as a dom. You may edit their orders at {FLASK_URL}") class DomRemoveCommand(TelegramCommand): command_text = "/dom_remove" command_regex = re.compile(r"^\/dom_remove( (?P@?\w+))$") description = "Remove a dom from your profile" async def exec_inner(self, text, update, session, *args, **kwargs): dom_username = text.split(' ')[1] if dom_username.startswith('@'): dom_username = dom_username[1:] sub_username = update['message']['chat']['username'] sub = user_get(sub_username) try: dom = user_get(dom_username) except: yield "I'm sorry that user isn't registered with me." return try: domsubusers_delete(sub, dom) except peewee.DoesNotExist: yield "I'm sorry that user is not on your list of doms" return yield f"Successfully removed {dom_username} from your list of doms" t = Telegram(session) await t.message_send(dom.telegram_chat_id, f"@{sub_username} has removed you as a dom.") class SkipDayAddCommand(TelegramCommand): command_text = "/skip_day" description = "Add a skip day" async def exec_inner(self, text, update, session, forResponse=None, reply=None): try: yield "Please enter a skip day" skip_day_str = await forResponse() try: skip_day = datetime.date.fromisoformat(skip_day_str) except: yield "The date you entered was not valid" return user = user_get(update['message']['chat']['username']) inline_keyboard = [{ 'text': 'All', 'callback_data': f"{CALLBACK_PREFIX_SKIP_DAY_ADD}_all" }] + [ { 'text': pool.name, 'callback_data': f"{CALLBACK_PREFIX_SKIP_DAY_ADD}_{pool.id}" } for pool in user.orders_pools ] yield "Please choose an order pool", { "inline_keyboard": [inline_keyboard] } pool_str = await forResponse() if not pool_str.startswith("skip_day_add_"): yield "There was a problem with your command" return try: order_pool_id = int(pool_str[len(CALLBACK_PREFIX_SKIP_DAY_ADD)+1:]) order_pool = user.orders_pools.where(OrdersPool.id == order_pool_id).get() except: order_pool = None skip_day_put(user, skip_day, order_pool) if order_pool is not None: yield f"Skip day {skip_day_str} for pool {order_pool.name} has been added" else: yield f"Skip day {skip_day_str} has been added" except ValueError: yield "That date was not valid" except peewee.IntegrityError: yield "That day has already been added" class SkipDayDeleteCommand(TelegramCommand): command_text = "/skip_day_delete" description = "Delete a skip day" async def exec_inner(self, text, update, session, forResponse=None, reply=None): try: yield "Please enter a skip day to delete" response = await forResponse() user = user_get(update['message']['chat']['username']) skip_day = datetime.date.fromisoformat(response) rows = skip_day_delete(user, skip_day) if rows > 0: yield f"Skip day {skip_day} has been deleted" else: yield f"Skip day {skip_day} does not exist" except ValueError: yield "That date was not valid" class SkipDaysCommand(TelegramCommand): command_text = "/skip_days" description = "List upcoming skip days" async def exec_inner(self, text, update, session, forResponse=None, reply=None): user = user_get(update['message']['chat']['username']) dates = skip_days_upcoming(user) response = "Upcoming skip days -\n" for date in dates: if date.pool is None: response += date.date.isoformat() + "\n" else: response += f"{date.date.isoformat()} - {date.pool.name}\n" yield response commands = [ StartCommand(), DomAddCommand(), DomRemoveCommand(), SkipDayAddCommand(), SkipDayDeleteCommand(), SkipDaysCommand() ]