From 2f9d7a708e8a1014204551d5ca670d29859ce6ac Mon Sep 17 00:00:00 2001 From: Frank Becker Date: Thu, 6 Feb 2020 20:56:53 +0100 Subject: [PATCH] add stop and start message to BotModule class. add remove_callback --- bot.py | 72 ++++++++++++++++---------------- modules/bot.py | 5 ++- modules/common/module.py | 7 ++-- modules/common/pollingservice.py | 3 +- modules/googlecal.py | 13 +++--- modules/help.py | 6 +-- modules/ig.py | 4 +- modules/loc.py | 5 +++ modules/twitter.py | 4 +- modules/url.py | 26 ++++++++---- 10 files changed, 81 insertions(+), 64 deletions(-) diff --git a/bot.py b/bot.py index 6796ce1..639a777 100755 --- a/bot.py +++ b/bot.py @@ -53,6 +53,12 @@ class Bot: } await self.client.room_send(room.room_id, 'm.room.message', msg) + def remove_callback(self, callback): + for cb_object in bot.client.event_callbacks: + if cb_object.func == callback: + print("remove callback") + bot.client.event_callbacks.remove(cb_object) + def get_room_by_id(self, room_id): return self.client.rooms[room_id] @@ -82,11 +88,10 @@ class Bot: def save_settings(self): module_settings = dict() for modulename, moduleobject in self.modules.items(): - if "get_settings" in dir(moduleobject): - try: - module_settings[modulename] = moduleobject.get_settings() - except Exception: - traceback.print_exc(file=sys.stderr) + try: + module_settings[modulename] = moduleobject.get_settings() + except Exception: + traceback.print_exc(file=sys.stderr) data = {self.appid: self.version, 'module_settings': module_settings} self.set_account_data(data) @@ -97,12 +102,11 @@ class Bot: return for modulename, moduleobject in self.modules.items(): if data['module_settings'].get(modulename): - if "set_settings" in dir(moduleobject): - try: - moduleobject.set_settings( - data['module_settings'][modulename]) - except Exception: - traceback.print_exc(file=sys.stderr) + try: + moduleobject.set_settings( + data['module_settings'][modulename]) + except Exception: + traceback.print_exc(file=sys.stderr) async def message_cb(self, room, event): # Figure out the command @@ -119,7 +123,7 @@ class Bot: moduleobject = self.modules.get(command) - if moduleobject.enabled and ("matrix_message" in dir(moduleobject)): + if moduleobject.enabled: try: await moduleobject.matrix_message(bot, room, event) except CommandRequiresAdmin: @@ -154,7 +158,7 @@ class Bot: module = importlib.import_module('modules.' + modulename) module = reload(module) cls = getattr(module, 'MatrixModule') - return cls() + return cls(modulename) except ModuleNotFoundError: print('Module ', modulename, ' failed to load!') traceback.print_exc(file=sys.stderr) @@ -184,11 +188,10 @@ class Bot: self.pollcount = self.pollcount + 1 for modulename, moduleobject in self.modules.items(): if moduleobject.enabled: - if "matrix_poll" in dir(moduleobject): - try: - await moduleobject.matrix_poll(bot, self.pollcount) - except Exception: - traceback.print_exc(file=sys.stderr) + try: + await moduleobject.matrix_poll(bot, self.pollcount) + except Exception: + traceback.print_exc(file=sys.stderr) await asyncio.sleep(10) def set_account_data(self, data): @@ -222,15 +225,14 @@ class Bot: print("NOTE: check MATRIX_ACCESS_TOKEN or set MATRIX_PASSWORD") sys.exit(2) - def init(self): - self.matrix_user = os.getenv('MATRIX_USER') - self.matrix_pass = os.getenv('MATRIX_PASSWORD') - matrix_server = os.getenv('MATRIX_SERVER') - bot_owners = os.getenv('BOT_OWNERS') - access_token = os.getenv('MATRIX_ACCESS_TOKEN') - join_on_invite = os.getenv('JOIN_ON_INVITE') + self.matrix_user = os.getenv('MATRIX_USER') + self.matrix_pass = os.getenv('MATRIX_PASSWORD') + matrix_server = os.getenv('MATRIX_SERVER') + bot_owners = os.getenv('BOT_OWNERS') + access_token = os.getenv('MATRIX_ACCESS_TOKEN') + join_on_invite = os.getenv('JOIN_ON_INVITE') if matrix_server and self.matrix_user and bot_owners: self.client = AsyncClient(matrix_server, self.matrix_user) @@ -256,22 +258,18 @@ class Bot: print(f'Starting {len(enabled_modules)} modules..') for modulename, moduleobject in self.modules.items(): if moduleobject.enabled: - print('Starting', modulename, '..') - if "matrix_start" in dir(moduleobject): - try: - moduleobject.matrix_start(bot) - except Exception: - traceback.print_exc(file=sys.stderr) + try: + moduleobject.matrix_start(bot) + except Exception: + traceback.print_exc(file=sys.stderr) def stop(self): print(f'Stopping {len(self.modules)} modules..') for modulename, moduleobject in self.modules.items(): - print('Stopping', modulename, '..') - if "matrix_stop" in dir(moduleobject): - try: - moduleobject.matrix_stop(bot) - except Exception: - traceback.print_exc(file=sys.stderr) + try: + moduleobject.matrix_stop(bot) + except Exception: + traceback.print_exc(file=sys.stderr) async def run(self): if not self.client.access_token: diff --git a/modules/bot.py b/modules/bot.py index 39fc2a5..eaaf5bf 100644 --- a/modules/bot.py +++ b/modules/bot.py @@ -4,11 +4,12 @@ from modules.common.module import BotModule class MatrixModule(BotModule): - def __init__(self): - super().__init__() + def __init__(self, name): + super().__init__(name) self.enable() def matrix_start(self, bot): + super().matrix_start(bot) self.starttime = datetime.now() async def matrix_message(self, bot, room, event): diff --git a/modules/common/module.py b/modules/common/module.py index aaa62bf..db39f2e 100644 --- a/modules/common/module.py +++ b/modules/common/module.py @@ -25,8 +25,9 @@ class BotModule(ABC): """ - def __init__(self): + def __init__(self, name): self.enabled = False + self.name = name def matrix_start(self, bot): """Called once on startup @@ -34,7 +35,7 @@ class BotModule(ABC): :param bot: a reference to the bot :type bot: Bot """ - pass + print('Starting', self.name, '..') @abstractmethod async def matrix_message(self, bot, room, event): @@ -55,7 +56,7 @@ class BotModule(ABC): :param bot: a reference to the bot :type bot: Bot """ - pass + print('Stopping', self.name, '..') async def matrix_poll(self, bot, pollcount): """Called every 10 seconds diff --git a/modules/common/pollingservice.py b/modules/common/pollingservice.py index 88623a0..bedac51 100644 --- a/modules/common/pollingservice.py +++ b/modules/common/pollingservice.py @@ -5,7 +5,8 @@ from modules.common.module import BotModule class PollingService(BotModule): - def __init__(self): + def __init__(self, name): + super().__init__(name) self.known_ids = set() self.account_rooms = dict() # Roomid -> [account, account..] self.next_poll_time = dict() # Roomid -> datetime, None = not polled yet diff --git a/modules/googlecal.py b/modules/googlecal.py index 0a49cfd..ccfb6c1 100644 --- a/modules/googlecal.py +++ b/modules/googlecal.py @@ -21,13 +21,17 @@ from modules.common.module import BotModule class MatrixModule(BotModule): - def matrix_start(self, bot): - self.bot = bot - self.SCOPES = ['https://www.googleapis.com/auth/calendar.readonly'] + def __init__(self, name): + super().__init__(name) self.credentials_file = "credentials.json" + self.SCOPES = ['https://www.googleapis.com/auth/calendar.readonly'] + self.bot = None self.service = None self.calendar_rooms = dict() # Contains room_id -> [calid, calid] .. + def matrix_start(self, bot): + super().matrix_start(bot) + self.bot = bot creds = None if not os.path.exists(self.credentials_file) or os.path.getsize(self.credentials_file) == 0: @@ -43,8 +47,7 @@ class MatrixModule(BotModule): if creds and creds.expired and creds.refresh_token: creds.refresh(Request()) else: - flow = InstalledAppFlow.from_client_secrets_file( - self.credentials_file, self.SCOPES) + flow = InstalledAppFlow.from_client_secrets_file(self.credentials_file, self.SCOPES) # urn:ietf:wg:oauth:2.0:oob creds = flow.run_local_server(port=0) # Save the credentials for the next run diff --git a/modules/help.py b/modules/help.py index dc77952..2914165 100644 --- a/modules/help.py +++ b/modules/help.py @@ -3,8 +3,8 @@ from modules.common.module import BotModule class MatrixModule(BotModule): - def __init__(self): - super().__init__() + def __init__(self, name): + super().__init__(name) self.enable() async def matrix_message(self, bot, room, event): @@ -17,7 +17,7 @@ class MatrixModule(BotModule): msg = msg + ' - ' + moduleobject.help() + '\n' except AttributeError: pass - msg + msg + '\n' + msg = msg + '\n' msg = msg + "\nAdd your own commands at https://github.com/vranki/hemppa" await bot.send_text(room, msg) diff --git a/modules/ig.py b/modules/ig.py index 15a4226..7ec3e46 100644 --- a/modules/ig.py +++ b/modules/ig.py @@ -11,8 +11,8 @@ from modules.common.pollingservice import PollingService class MatrixModule(PollingService): - def __init__(self): - super().__init__() + def __init__(self, name): + super().__init__(name) self.instagram = Instagram() self.service_name = 'Instagram' diff --git a/modules/loc.py b/modules/loc.py index 7e1217f..8f0c475 100644 --- a/modules/loc.py +++ b/modules/loc.py @@ -7,9 +7,14 @@ class MatrixModule(BotModule): bot = None def matrix_start(self, bot): + super().matrix_start(bot) self.bot = bot bot.client.add_event_callback(self.unknown_cb, RoomMessageUnknown) + def matrix_stop(self, bot): + super().matrix_stop(bot) + bot.remove_callback(self.unknown_cb) + async def unknown_cb(self, room, event): if event.msgtype != 'm.location': return diff --git a/modules/twitter.py b/modules/twitter.py index 4479368..0582ea7 100644 --- a/modules/twitter.py +++ b/modules/twitter.py @@ -9,8 +9,8 @@ from modules.common.pollingservice import PollingService # https://github.com/taspinar/twitterscraper/tree/master/twitterscraper class MatrixModule(PollingService): - def __init__(self): - super().__init__() + def __init__(self, name): + super().__init__(name) self.service_name = 'Twitter' async def poll_implementation(self, bot, account, roomid, send_messages): diff --git a/modules/url.py b/modules/url.py index 3f37e56..3536d99 100644 --- a/modules/url.py +++ b/modules/url.py @@ -4,7 +4,7 @@ from functools import lru_cache import httpx from bs4 import BeautifulSoup -from nio import RoomMessageText +from nio import RoomMessageText, AsyncClient from modules.common.module import BotModule @@ -16,23 +16,31 @@ class MatrixModule(BotModule): Everytime a url is seen in a message we do http request to it and try to get a title tag contents to spit out to the room. """ - bot = None - status = dict() # room_id -> what to do with urls + def __init__(self, name): + super().__init__(name) - STATUSES = { - "OFF": "Not spamming this channel", - "TITLE": "Spamming this channel with titles", - "DESCRIPTION": "Spamming this channel with descriptions", - "BOTH": "Spamming this channel with both title and description", - } + self.bot = None + self.status = dict() # room_id -> what to do with urls + + self.STATUSES = { + "OFF": "Not spamming this channel", + "TITLE": "Spamming this channel with titles", + "DESCRIPTION": "Spamming this channel with descriptions", + "BOTH": "Spamming this channel with both title and description", + } def matrix_start(self, bot): """ Register callback for all RoomMessageText events on startup """ + super().matrix_start(bot) self.bot = bot bot.client.add_event_callback(self.text_cb, RoomMessageText) + def matrix_stop(self, bot): + super().matrix_stop(bot) + bot.remove_callback(self.text_cb) + async def text_cb(self, room, event): """ Handle client callbacks for all room text events