gear-orders/telegram/telegram.py

import logging
import asyncio

from settings import ENV, TELEGRAM_API_TOKEN, TELEGRAM_ALLOWLIST, TELEGRAM_COMMAND_TIMEOUT
from util import make_session

logger = logging.getLogger(__name__)

TELEGRAM_API_URL = 'https://api.telegram.org/bot'

class Telegram:
  def __init__(self, aiosession):
    if TELEGRAM_API_TOKEN is None:
      raise Exception("TELEGRAM_API_TOKEN is required")
    
    self.aiosession = aiosession

  async def get(self, request, params):
    async with self.aiosession.get('{}{}/{}'.format(
      TELEGRAM_API_URL,
      TELEGRAM_API_TOKEN,
      request
    ), params=params) as response:
      if not response.ok:
        logger.error('{} {}'.format(request, response.status))
        raise Exception('{} {}'.format(request, response.status))

      return await response.json()

  async def post(self, request, data):
    async with self.aiosession.post('{}{}/{}'.format(
      TELEGRAM_API_URL,
      TELEGRAM_API_TOKEN,
      request
    ), json=data) as response:
      if not response.ok:
        logger.error('{} {} {}'.format(request, response.status, await response.json()))
        raise Exception('{} {}'.format(request, response.status))
      
      return await response.json()
    
  async def message_send(self, chat_id, text, reply_markup=None):
    data = {
      'chat_id': chat_id,
      'text': text
    }
    
    if reply_markup is not None:
      data['reply_markup'] = reply_markup

    return await self.post('sendMessage', data=data)
  
  async def message_edit_reply_markup(self, chat_id, message_id, reply_markup):
    return await self.post('editMessageReplyMarkup', data={
      'chat_id': chat_id,
      'message_id': message_id,
      'reply_markup': reply_markup
    })
  
  async def set_commands(self, commands):
    data = {
      'commands': commands
    }

    return await self.post('setMyCommands', data=data)
  
  async def updates_get(self, offset=0, timeout=240):
    return await self.get('getUpdates', params={
      'offset': offset,
      'timeout': timeout
    })

class TelegramCommand():
  command_regex = None
  command_text = None
  authorized = True

  def matches(self, text):
    if self.command_regex is not None:
      return self.command_regex.match(text) is not None
    
    if self.command_text is not None:
      return text == self.command_text
    
    return False

  async def exec(self, *args, **kwargs):
    try:
      async for reply in self.exec_inner(*args, **kwargs):
        yield reply
    except Exception as e:
      yield "There was a problem with your command"
      logger.error(f"{self.__class__.__name__} {e}")

  async def exec_inner(self, text, update, session):
    pass

async def handle_commands(commands=[], loop=None):
  async with make_session() as session:
    t = Telegram(session)

    await t.set_commands([
      {
        "command": c.command_text,
        "description": c.description
      } for c in commands
      if c.command_text is not None and c.description is not None
    ])

    command_tasks = set()
    command_futures = {}
    def forResponse():
      f = loop.create_future()
      command_futures[chat_id] = f
      return f

    offset = 0
    while True:
      try:
        updates = await t.updates_get(offset=offset)

        if(updates['ok']):
          for update in updates['result']:
            if('message' in update):
              chat_id = update['message']['chat']['id']

              if(update['message']['chat']['type'] != "private" or
                 update['message']['chat']['username'] not in TELEGRAM_ALLOWLIST):
                continue

              if('text' not in update['message']):
                continue

              if(chat_id in command_futures):
                command_futures[chat_id].set_result(
                  update['message']['text']
                )
                del command_futures[chat_id]
                continue

              for command in commands:
                if(command.matches(update['message']['text'])):
                  async def timeoutTask():
                    try:
                      async with asyncio.timeout(TELEGRAM_COMMAND_TIMEOUT):
                          async for reply in command.exec(
                              update['message']['text'],
                              update,
                              session,
                              forResponse=forResponse
                            ):
                            match reply.__class__.__name__:
                              case 'str':
                                await t.message_send(chat_id, reply)
                              case 'tuple':
                                (text, reply_markup) = reply
                                await t.message_send(chat_id, text, reply_markup=reply_markup)
                    except TimeoutError:
                      if chat_id in command_futures:
                        del command_futures[chat_id]
                      
                      await t.message_send(chat_id, "Your command has timed out")
                    except Exception:
                      await t.message_send(chat_id, 'There was a problem with your command')
                      logger.exception('Problem while executing a command')
                    
                  task = asyncio.create_task(
                    timeoutTask()
                  )

                  command_tasks.add(task)
                  task.add_done_callback(command_tasks.discard)

                  break
            elif('callback_query' in update):
              chat_id = update['callback_query']['message']['chat']['id']

              if(update['callback_query']['message']['chat']['type'] != "private" or
                 update['callback_query']['message']['chat']['username'] not in TELEGRAM_ALLOWLIST):
                continue

              if(chat_id in command_futures):
                command_futures[chat_id].set_result(
                  update['callback_query']['data']
                )
                del command_futures[chat_id]

                # Clear inline keyboard
                await t.message_edit_reply_markup(
                  chat_id,
                  update['callback_query']['message']['message_id'],
                  { "inline_keyboard": [] }
                )
                continue

            offset = update['update_id'] + 1
        else:
          logger.warning("getUpdates was not ok {}".format(update))
          await asyncio.sleep(30)
      except Exception as e:
        logger.exception("getUpdates exception {}".format(e))
        await asyncio.sleep(30)