Cleanup init and owner setup

This commit is contained in:
2022-06-09 11:45:46 +01:00
parent 5a61444535
commit db04de242a
2 changed files with 13 additions and 6 deletions

View File

@@ -28,6 +28,8 @@ def main():
# Set the owner ID if configured # Set the owner ID if configured
if 'SMSBOT_OWNER_ID' in os.environ: if 'SMSBOT_OWNER_ID' in os.environ:
telegram_bot.set_owner(os.environ.get('SMSBOT_OWNER_ID')) telegram_bot.set_owner(os.environ.get('SMSBOT_OWNER_ID'))
else:
logging.warning('No Owner ID is set, which is not a good idea...')
# Add default subscribers # Add default subscribers
if 'SMSBOT_DEFAULT_SUBSCRIBERS' in os.environ: if 'SMSBOT_DEFAULT_SUBSCRIBERS' in os.environ:

View File

@@ -8,21 +8,24 @@ from smsbot.utils import get_smsbot_version
class TelegramSmsBot(object): class TelegramSmsBot(object):
def __init__(self, token, owner=None, subscribers=None): def __init__(self, token, allow_subscribing=False, owner=None, subscribers=None):
self.logger = logging.getLogger(self.__class__.__name__) self.logger = logging.getLogger(self.__class__.__name__)
self.bot_token = token self.bot_token = token
self.owner_id = owner
self.subscriber_ids = subscribers or [] self.subscriber_ids = subscribers or []
self.set_owner(owner)
def start(self):
self.logger.info('Starting bot...')
self.updater = Updater(self.bot_token, use_context=True) self.updater = Updater(self.bot_token, use_context=True)
self.updater.dispatcher.add_handler(CommandHandler('help', self.help_handler)) self.updater.dispatcher.add_handler(CommandHandler('help', self.help_handler))
self.updater.dispatcher.add_handler(CommandHandler('start', self.help_handler)) self.updater.dispatcher.add_handler(CommandHandler('start', self.help_handler))
self.updater.dispatcher.add_handler(CommandHandler('subscribe', self.subscribe_handler))
self.updater.dispatcher.add_handler(CommandHandler('unsubscribe', self.unsubscribe_handler)) if allow_subscribing:
self.updater.dispatcher.add_handler(CommandHandler('subscribe', self.subscribe_handler))
self.updater.dispatcher.add_handler(CommandHandler('unsubscribe', self.unsubscribe_handler))
self.updater.dispatcher.add_error_handler(self.error_handler) self.updater.dispatcher.add_error_handler(self.error_handler)
def start(self):
self.logger.info('Starting bot...')
self.updater.start_polling() self.updater.start_polling()
self.bot = self.updater.bot self.bot = self.updater.bot
self.logger.info('Bot Ready') self.logger.info('Bot Ready')
@@ -78,6 +81,8 @@ class TelegramSmsBot(object):
def set_owner(self, chat_id): def set_owner(self, chat_id):
self.owner_id = chat_id self.owner_id = chat_id
if self.owner_id and self.owner_id not in self.subscriber_ids:
self.subscriber_ids.append(self.owner_id)
def add_subscriber(self, chat_id): def add_subscriber(self, chat_id):
self.subscriber_ids.append(chat_id) self.subscriber_ids.append(chat_id)