hemppa/bot.py

251 lines
8.7 KiB
Python
Executable File

#!/usr/bin/env python3
import asyncio
import glob
import importlib
import json
import os
import re
import sys
import traceback
import urllib.parse
import requests
from nio import AsyncClient, InviteEvent, JoinError, RoomMessageText
# 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,
# or is bot owner
def is_admin(self, room, event):
if self.is_owner(event):
return True
if event.sender not 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 Exception:
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 Exception:
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 Exception:
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(
f'Received invite event, but not joining as sender is not owner or bot not configured to join on invite. {event}')
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 Exception:
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 Exception:
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 Exception:
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()