#!/usr/bin/env python3 import asyncio import os import json import glob import traceback import importlib import sys import re import requests import json import urllib.parse from nio import (AsyncClient, RoomMessageText, RoomMessageUnknown, JoinError, InviteEvent) # Couple of custom exceptions class CommandRequiresAdmin(Exception): pass class CommandRequiresOwner(Exception): pass class Bot: appid = 'org.vranki.hemppa' version = '1.1' client = None join_on_invite = False modules = dict() pollcount = 0 poll_task = None owners = [] async def send_text(self, room, body): msg = { "body": body, "msgtype": "m.text" } await self.client.room_send(room.room_id, 'm.room.message', msg) async def send_html(self, room, html, plaintext): msg = { "msgtype": "m.text", "format": "org.matrix.custom.html", "formatted_body": html, "body": plaintext } await self.client.room_send(room.room_id, 'm.room.message', msg) def get_room_by_id(self, room_id): return self.client.rooms[room_id] # Throws exception if event sender is not a room admin def must_be_admin(self, room, event): if not self.is_admin(room, event): raise CommandRequiresAdmin # Throws exception if event sender is not a bot owner def must_be_owner(self, event): if not self.is_owner(event): raise CommandRequiresOwner # Returns true if event's sender is admin in the room event was sent in def is_admin(self, room, event): if not event.sender in room.power_levels.users: return False return room.power_levels.users[event.sender] >= 50 # Returns true if event's sender is owner of the bot def is_owner(self, event): return event.sender in self.owners 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: traceback.print_exc(file=sys.stderr) data = { self.appid: self.version, 'module_settings': module_settings} self.set_account_data(data) def load_settings(self, data): if not data: return if not data.get('module_settings'): 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: traceback.print_exc(file=sys.stderr) async def message_cb(self, room, event): # Figure out the command body = event.body if len(body) == 0: return command = body.split().pop(0) # Strip away non-alphanumeric characters, including leading ! for security command = re.sub(r'\W+', '', command) moduleobject = self.modules.get(command) if "matrix_message" in dir(moduleobject): try: await moduleobject.matrix_message(bot, room, event) except CommandRequiresAdmin: await self.send_text(room, f'Sorry, you need admin power level in this room to run that command.') except CommandRequiresOwner: await self.send_text(room, f'Sorry, only bot owner can run that command.') except: await self.send_text(room, f'Module {command} experienced difficulty: {sys.exc_info()[0]} - see log for details') traceback.print_exc(file=sys.stderr) async def invite_cb(self, room, event): if self.join_on_invite or self.is_owner(event): for attempt in range(3): result = await self.client.join(room.room_id) if type(result) == JoinError: print(f"Error joining room {room.room_id} (attempt %d): %s", attempt, result.message, ) else: break else: print('Received invite event, but not joining as sender is not owner or bot not configured to join on invte.') def load_module(self, modulename): try: module = importlib.import_module('modules.' + modulename) cls = getattr(module, 'MatrixModule') return cls() except ModuleNotFoundError: print('Module ', modulename, ' failed to load!') traceback.print_exc(file=sys.stderr) return None def get_modules(self): modulefiles = glob.glob('./modules/*.py') for modulefile in modulefiles: modulename = os.path.splitext(os.path.basename(modulefile))[0] moduleobject = self.load_module(modulename) if moduleobject: self.modules[modulename] = moduleobject async def poll_timer(self): while True: self.pollcount = self.pollcount + 1 for modulename, moduleobject in self.modules.items(): if "matrix_poll" in dir(moduleobject): try: await moduleobject.matrix_poll(bot, self.pollcount) except: traceback.print_exc(file=sys.stderr) await asyncio.sleep(10) def set_account_data(self, data): userid = urllib.parse.quote(os.environ['MATRIX_USER']) ad_url = f"{self.client.homeserver}/_matrix/client/r0/user/{userid}/account_data/{self.appid}?access_token={self.client.access_token}" response = requests.put(ad_url, json.dumps(data)) if response.status_code != 200: print('Setting account data failed:', response, response.json()) def get_account_data(self): userid = urllib.parse.quote(os.environ['MATRIX_USER']) ad_url = f"{self.client.homeserver}/_matrix/client/r0/user/{userid}/account_data/{self.appid}?access_token={self.client.access_token}" response = requests.get(ad_url) if response.status_code == 200: return response.json() print(f'Getting account data failed: {response} {response.json()} - this is normal if you have not saved any settings yet.') return None def init(self): self.client = AsyncClient(os.environ['MATRIX_SERVER'], os.environ['MATRIX_USER']) self.client.access_token = os.getenv('MATRIX_ACCESS_TOKEN') self.join_on_invite = os.getenv("JOIN_ON_INVITE") is not None self.owners = os.environ['BOT_OWNERS'].split(',') self.get_modules() 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: traceback.print_exc(file=sys.stderr) async def run(self): if not self.client.access_token: await self.client.login(os.environ['MATRIX_PASSWORD']) print("Logged in with password, access token:", self.client.access_token) await self.client.sync() print(f'Starting {len(self.modules)} modules..') for modulename, moduleobject in self.modules.items(): print('Starting', modulename, '..') if "matrix_start" in dir(moduleobject): try: moduleobject.matrix_start(bot) except: traceback.print_exc(file=sys.stderr) self.poll_task = asyncio.get_event_loop().create_task(self.poll_timer()) if self.client.logged_in: self.load_settings(self.get_account_data()) self.client.add_event_callback(self.message_cb, RoomMessageText) self.client.add_event_callback(self.invite_cb, (InviteEvent,)) if self.join_on_invite: print('Note: Bot will join rooms if invited') print('Bot running as', self.client.user, ', owners', self.owners) await self.client.sync_forever(timeout=30000) else: print('Client was not able to log in, check env variables!') bot = Bot() bot.init() try: asyncio.get_event_loop().run_until_complete(bot.run()) except KeyboardInterrupt: if bot.poll_task: bot.poll_task.cancel() bot.stop()