diff options
Diffstat (limited to 'include/py/homekit/telegram/bot.py')
| -rw-r--r-- | include/py/homekit/telegram/bot.py | 574 |
1 files changed, 574 insertions, 0 deletions
diff --git a/include/py/homekit/telegram/bot.py b/include/py/homekit/telegram/bot.py new file mode 100644 index 0000000..f5f620a --- /dev/null +++ b/include/py/homekit/telegram/bot.py @@ -0,0 +1,574 @@ +from __future__ import annotations + +import logging +import itertools + +from enum import Enum, auto +from functools import wraps +from typing import Optional, Union, Tuple, Coroutine + +from telegram import Update, ReplyKeyboardMarkup +from telegram.ext import ( + Application, + filters, + CommandHandler, + MessageHandler, + CallbackQueryHandler, + CallbackContext, + ConversationHandler +) +from telegram.ext.filters import BaseFilter +from telegram.error import TimedOut + +from homekit.config import config + +from ._botlang import lang, languages +from ._botdb import BotDatabase +from ._botutil import exc2text, IgnoreMarkup +from ._botcontext import Context +from .config import TelegramUserListType + + +db: Optional[BotDatabase] = None + +_user_filter: Optional[BaseFilter] = None +_cancel_filter = filters.Text(lang.all('cancel')) +_back_filter = filters.Text(lang.all('back')) +_cancel_and_back_filter = filters.Text(lang.all('back') + lang.all('cancel')) + +_logger = logging.getLogger(__name__) +_application: Optional[Application] = None +_exception_handler: Optional[Coroutine] = None +_dispatcher = None +_markup_getter: Optional[callable] = None +_start_handler_ref: Optional[Coroutine] = None + + +def text_filter(*args): + if not _user_filter: + raise RuntimeError('user_filter is not initialized') + return filters.Text(args[0] if isinstance(args[0], list) else [*args]) & _user_filter + + +async def _handler_of_handler(*args, **kwargs): + self = None + context = None + update = None + + _args = list(args) + while len(_args): + v = _args[0] + if isinstance(v, conversation): + self = v + _args.pop(0) + elif isinstance(v, Update): + update = v + _args.pop(0) + elif isinstance(v, CallbackContext): + context = v + _args.pop(0) + break + + ctx = Context(update, + callback_context=context, + markup_getter=lambda _ctx: None if not _markup_getter else _markup_getter(_ctx), + store=db) + try: + _args.insert(0, ctx) + + f = kwargs['f'] + del kwargs['f'] + + if 'return_with_context' in kwargs: + return_with_context = True + del kwargs['return_with_context'] + else: + return_with_context = False + + if 'argument' in kwargs and kwargs['argument'] == 'message_key': + del kwargs['argument'] + mkey = None + for k, v in lang.get_langpack(ctx.user_lang).items(): + if ctx.text == v: + mkey = k + break + _args.insert(0, mkey) + + if self: + _args.insert(0, self) + + result = await f(*_args, **kwargs) + return result if not return_with_context else (result, ctx) + + except Exception as e: + if _exception_handler: + if not _exception_handler(e, ctx) and not isinstance(e, TimedOut): + _logger.exception(e) + if not ctx.is_callback_context(): + await ctx.reply_exc(e) + else: + notify_user(ctx.user_id, exc2text(e)) + else: + _logger.exception(e) + + +def handler(**kwargs): + def inner(f): + @wraps(f) + async def _handler(*args, **inner_kwargs): + if 'argument' in kwargs and kwargs['argument'] == 'message_key': + inner_kwargs['argument'] = 'message_key' + return await _handler_of_handler(f=f, *args, **inner_kwargs) + + messages = [] + texts = [] + + if 'messages' in kwargs: + messages += kwargs['messages'] + if 'message' in kwargs: + messages.append(kwargs['message']) + + if 'text' in kwargs: + texts.append(kwargs['text']) + if 'texts' in kwargs: + texts += kwargs['texts'] + + if messages or texts: + new_messages = list(itertools.chain.from_iterable([lang.all(m) for m in messages])) + texts += new_messages + texts = list(set(texts)) + _application.add_handler( + MessageHandler(text_filter(*texts), _handler), + group=0 + ) + + if 'command' in kwargs: + _application.add_handler(CommandHandler(kwargs['command'], _handler), group=0) + + if 'callback' in kwargs: + _application.add_handler(CallbackQueryHandler(_handler, pattern=kwargs['callback']), group=0) + + return _handler + + return inner + + +def simplehandler(f: Coroutine): + @wraps(f) + async def _handler(*args, **kwargs): + return await _handler_of_handler(f=f, *args, **kwargs) + return _handler + + +def callbackhandler(*args, **kwargs): + def inner(f): + @wraps(f) + async def _handler(*args, **kwargs): + return await _handler_of_handler(f=f, *args, **kwargs) + pattern_kwargs = {} + if kwargs['callback'] != '*': + pattern_kwargs['pattern'] = kwargs['callback'] + _application.add_handler(CallbackQueryHandler(_handler, **pattern_kwargs), group=0) + return _handler + return inner + + +async def exceptionhandler(f: callable): + global _exception_handler + if _exception_handler: + _logger.warning('exception handler already set, we will overwrite it') + _exception_handler = f + + +def defaultreplymarkup(f: callable): + global _markup_getter + _markup_getter = f + + +def convinput(state, is_enter=False, **kwargs): + def inner(f): + f.__dict__['_conv_data'] = dict( + orig_f=f, + enter=is_enter, + type=ConversationMethodType.ENTRY if is_enter and state == 0 else ConversationMethodType.STATE_HANDLER, + state=state, + **kwargs + ) + + @wraps(f) + async def _impl(*args, **kwargs): + result, ctx = await _handler_of_handler(f=f, *args, **kwargs, return_with_context=True) + if result == conversation.END: + await start(ctx) + return result + + return _impl + + return inner + + +def conventer(state, **kwargs): + return convinput(state, is_enter=True, **kwargs) + + +class ConversationMethodType(Enum): + ENTRY = auto() + STATE_HANDLER = auto() + + +class conversation: + END = ConversationHandler.END + STATE_SEQS = [] + + def __init__(self, enable_back=False): + self._logger = logging.getLogger(self.__class__.__name__) + self._user_state_cache = {} + self._back_enabled = enable_back + + def make_handlers(self, f: callable, **kwargs) -> list: + messages = {} + handlers = [] + + if 'messages' in kwargs: + if isinstance(kwargs['messages'], dict): + messages = kwargs['messages'] + else: + for m in kwargs['messages']: + messages[m] = None + + if 'message' in kwargs: + if isinstance(kwargs['message'], str): + messages[kwargs['message']] = None + else: + AttributeError('invalid message type: ' + type(kwargs['message'])) + + if messages: + for message, target_state in messages.items(): + if not target_state: + handlers.append(MessageHandler(text_filter(lang.all(message) if 'messages_lang_completed' not in kwargs else message), f)) + else: + handlers.append(MessageHandler(text_filter(lang.all(message) if 'messages_lang_completed' not in kwargs else message), self.make_invoker(target_state))) + + if 'regex' in kwargs: + handlers.append(MessageHandler(filters.Regex(kwargs['regex']) & _user_filter, f)) + + if 'command' in kwargs: + handlers.append(CommandHandler(kwargs['command'], f, _user_filter)) + + return handlers + + def make_invoker(self, state): + def _invoke(update: Update, context: CallbackContext): + ctx = Context(update, + callback_context=context, + markup_getter=lambda _ctx: None if not _markup_getter else _markup_getter(_ctx), + store=db) + return self.invoke(state, ctx) + return _invoke + + async def invoke(self, state, ctx: Context): + self._logger.debug(f'invoke, state={state}') + for item in dir(self): + f = getattr(self, item) + if not callable(f) or item.startswith('_') or '_conv_data' not in f.__dict__: + continue + cd = f.__dict__['_conv_data'] + if cd['enter'] and cd['state'] == state: + return await cd['orig_f'](self, ctx) + + raise RuntimeError(f'invoke: failed to find method for state {state}') + + def get_handler(self) -> ConversationHandler: + entry_points = [] + states = {} + + l_cancel_filter = _cancel_filter if not self._back_enabled else _cancel_and_back_filter + + for item in dir(self): + f = getattr(self, item) + if not callable(f) or item.startswith('_') or '_conv_data' not in f.__dict__: + continue + + cd = f.__dict__['_conv_data'] + + if cd['type'] == ConversationMethodType.ENTRY: + entry_points = self.make_handlers(f, **cd) + elif cd['type'] == ConversationMethodType.STATE_HANDLER: + states[cd['state']] = self.make_handlers(f, **cd) + states[cd['state']].append( + MessageHandler(_user_filter & ~l_cancel_filter, conversation.invalid) + ) + + fallbacks = [MessageHandler(_user_filter & _cancel_filter, self.cancel)] + if self._back_enabled: + fallbacks.append(MessageHandler(_user_filter & _back_filter, self.back)) + + return ConversationHandler( + entry_points=entry_points, + states=states, + fallbacks=fallbacks + ) + + def get_user_state(self, user_id: int) -> Optional[int]: + if user_id not in self._user_state_cache: + return None + return self._user_state_cache[user_id] + + # TODO store in ctx.user_state + def set_user_state(self, user_id: int, state: Union[int, None]): + if not self._back_enabled: + return + if state is not None: + self._user_state_cache[user_id] = state + else: + del self._user_state_cache[user_id] + + @staticmethod + @simplehandler + async def invalid(ctx: Context): + await ctx.reply(ctx.lang('invalid_input'), markup=IgnoreMarkup()) + # return 0 # FIXME is this needed + + @simplehandler + async def cancel(self, ctx: Context): + await start(ctx) + self.set_user_state(ctx.user_id, None) + return conversation.END + + @simplehandler + async def back(self, ctx: Context): + cur_state = self.get_user_state(ctx.user_id) + if cur_state is None: + await start(ctx) + self.set_user_state(ctx.user_id, None) + return conversation.END + + new_state = None + for seq in self.STATE_SEQS: + if cur_state in seq: + idx = seq.index(cur_state) + if idx > 0: + return self.invoke(seq[idx-1], ctx) + + if new_state is None: + raise RuntimeError('failed to determine state to go back to') + + @classmethod + def add_cancel_button(cls, ctx: Context, buttons): + buttons.append([ctx.lang('cancel')]) + + @classmethod + def add_back_button(cls, ctx: Context, buttons): + # buttons.insert(0, [ctx.lang('back')]) + buttons.append([ctx.lang('back')]) + + async def reply(self, + ctx: Context, + state: Union[int, Enum], + text: str, + buttons: Optional[list], + with_cancel=False, + with_back=False, + buttons_lang_completed=False): + + if buttons: + new_buttons = [] + if not buttons_lang_completed: + for item in buttons: + if isinstance(item, list): + item = map(lambda s: ctx.lang(s), item) + new_buttons.append(list(item)) + elif isinstance(item, str): + new_buttons.append([ctx.lang(item)]) + else: + raise ValueError('invalid type: ' + type(item)) + else: + new_buttons = list(buttons) + + buttons = None + else: + if with_cancel or with_back: + new_buttons = [] + else: + new_buttons = None + + if with_cancel: + self.add_cancel_button(ctx, new_buttons) + if with_back: + if not self._back_enabled: + raise AttributeError(f'back is not enabled for this conversation ({self.__class__.__name__})') + self.add_back_button(ctx, new_buttons) + + markup = ReplyKeyboardMarkup(new_buttons, one_time_keyboard=True) if new_buttons else IgnoreMarkup() + await ctx.reply(text, markup=markup) + self.set_user_state(ctx.user_id, state) + return state + + +class LangConversation(conversation): + START, = range(1) + + @conventer(START, command='lang') + async def entry(self, ctx: Context): + self._logger.debug(f'current language: {ctx.user_lang}') + + buttons = [] + for name in languages.values(): + buttons.append(name) + markup = ReplyKeyboardMarkup([buttons, [ctx.lang('cancel')]], one_time_keyboard=False) + + await ctx.reply(ctx.lang('select_language'), markup=markup) + return self.START + + @convinput(START, messages=lang.languages) + async def input(self, ctx: Context): + selected_lang = None + for key, value in languages.items(): + if value == ctx.text: + selected_lang = key + break + + if selected_lang is None: + raise ValueError('could not find the language') + + db.set_user_lang(ctx.user_id, selected_lang) + await ctx.reply(ctx.lang('saved'), markup=IgnoreMarkup()) + + return self.END + + +def initialize(): + global _user_filter + global _application + # global _updater + global _dispatcher + + # init user_filter + _user_ids = config.app_config.get_user_ids() + if len(_user_ids) > 0: + _logger.info('allowed users: ' + str(_user_ids)) + _user_filter = filters.User(_user_ids) + else: + _user_filter = filters.ALL # not sure if this is correct + + _application = Application.builder()\ + .token(config.app_config.get('bot.token'))\ + .connect_timeout(7)\ + .read_timeout(6)\ + .build() + + # transparently log all messages + # _application.dispatcher.add_handler(MessageHandler(filters.ALL & _user_filter, _logging_message_handler), group=10) + # _application.dispatcher.add_handler(CallbackQueryHandler(_logging_callback_handler), group=10) + + +def run(start_handler=None, any_handler=None): + global db + global _start_handler_ref + + if not start_handler: + start_handler = _default_start_handler + if not any_handler: + any_handler = _default_any_handler + if not db: + db = BotDatabase() + + _start_handler_ref = start_handler + + _application.add_handler(LangConversation().get_handler(), group=0) + _application.add_handler(CommandHandler('start', + callback=simplehandler(start_handler), + filters=_user_filter)) + _application.add_handler(MessageHandler(filters.ALL & _user_filter, any_handler)) + + _application.run_polling() + + +def add_conversation(conv: conversation) -> None: + _application.add_handler(conv.get_handler(), group=0) + + +def add_handler(h): + _application.add_handler(h, group=0) + + +async def start(ctx: Context): + return await _start_handler_ref(ctx) + + +async def _default_start_handler(ctx: Context): + if 'start_message' not in lang: + return await ctx.reply('Please define start_message or override start()') + await ctx.reply(ctx.lang('start_message')) + + +@simplehandler +async def _default_any_handler(ctx: Context): + if 'invalid_command' not in lang: + return await ctx.reply('Please define invalid_command or override any()') + await ctx.reply(ctx.lang('invalid_command')) + + +# def _logging_message_handler(update: Update, context: CallbackContext): +# if _reporting: +# _reporting.report(update.message) +# +# +# def _logging_callback_handler(update: Update, context: CallbackContext): +# if _reporting: +# _reporting.report(update.callback_query.message, text=update.callback_query.data) + + +async def notify_all(text_getter: callable, + exclude: Tuple[int] = ()) -> None: + notify_user_ids = config.app_config.get_user_ids(TelegramUserListType.NOTIFY) + if not notify_user_ids: + _logger.error('notify_all() called but no notify_users defined in the config') + return + + for user_id in notify_user_ids: + if user_id in exclude: + continue + + text = text_getter(db.get_user_lang(user_id)) + await _application.bot.send_message(chat_id=user_id, + text=text, + parse_mode='HTML') + + +async def notify_user(user_id: int, text: Union[str, Exception], **kwargs) -> None: + if isinstance(text, Exception): + text = exc2text(text) + await _application.bot.send_message(chat_id=user_id, + text=text, + parse_mode='HTML', + **kwargs) + + +async def send_photo(user_id, **kwargs): + await _application.bot.send_photo(chat_id=user_id, **kwargs) + + +async def send_audio(user_id, **kwargs): + await _application.bot.send_audio(chat_id=user_id, **kwargs) + + +async def send_file(user_id, **kwargs): + await _application.bot.send_document(chat_id=user_id, **kwargs) + + +async def edit_message_text(user_id, message_id, *args, **kwargs): + await _application.bot.edit_message_text(chat_id=user_id, + message_id=message_id, + parse_mode='HTML', + *args, **kwargs) + + +async def delete_message(user_id, message_id): + await _application.bot.delete_message(chat_id=user_id, message_id=message_id) + + +def set_database(_db: BotDatabase): + global db + db = _db + |