241 lines
9.1 KiB
Python
241 lines
9.1 KiB
Python
import logging
|
|
import argparse
|
|
import os
|
|
import asyncio
|
|
import yaml
|
|
import importlib
|
|
import pathlib
|
|
from nio import AsyncClient, LoginResponse
|
|
|
|
from skills import *
|
|
|
|
logging.basicConfig(format='[%(asctime)s] %(process)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s', level=logging.INFO)
|
|
|
|
class MatrixBot:
|
|
__AUTHOR__ = 'anima'
|
|
__VERSION__ = '0.6.3'
|
|
def __init__(self, config_file: str = 'bot.conf', session_file: str = 'session.tmp') -> None:
|
|
"""Initialize the MatrixBot instance.
|
|
|
|
Args:
|
|
config_file (str, optional): path to config file
|
|
session_file (str, optional): path to session file
|
|
"""
|
|
self.botname = 'Bot'
|
|
self.config_file = config_file
|
|
self.session_file = session_file
|
|
self._config = dict()
|
|
self.silence = False
|
|
self.skills = list()
|
|
self._find_skills()
|
|
|
|
if self._load_config():
|
|
self.client = AsyncClient(self.homeserver, self.username)
|
|
self.helper = BotHelper(self.client)
|
|
else:
|
|
logging.critical('no config file was found')
|
|
exit(1)
|
|
|
|
def _load_config(self) -> bool:
|
|
"""Load given config file.
|
|
|
|
Returns:
|
|
bool: True if config file was loaded successfully
|
|
"""
|
|
if not os.path.exists(self.config_file):
|
|
return False
|
|
logging.info(f'load config from file {self.config_file}')
|
|
with open(self.config_file, 'r') as f:
|
|
self._config = yaml.safe_load(f)
|
|
self.botname = self._config['general']['botname']
|
|
self.domain = self._config['general']['domain']
|
|
self.homeserver = f"https://{self.domain}"
|
|
self.username = f"@{self._config['general']['user']}:{self._config['general']['domain']}"
|
|
# get password from env or fallback from config file
|
|
self.__password = os.environ.get('MATRIX_BOT_PASSWORD') or self._config['general']['pass']
|
|
del self._config['general']['pass']
|
|
return True
|
|
|
|
def _save_session(self) -> bool:
|
|
"""Save session data as file.
|
|
|
|
Returns:
|
|
bool: True if file was saved
|
|
"""
|
|
try:
|
|
session = {
|
|
'access_token': self.session.access_token,
|
|
'device_id': self.session.device_id,
|
|
'next_batch': self.client.next_batch
|
|
}
|
|
logging.debug('save session file')
|
|
# Secure file handling: set appropriate permissions (owner read/write only)
|
|
old_umask = os.umask(0o077) # Restrict file permissions to owner only
|
|
try:
|
|
with open(self.session_file, 'w') as f:
|
|
yaml.dump(session, f)
|
|
# Explicitly set file permissions to owner-only read/write (no execute)
|
|
os.chmod(self.session_file, 0o600)
|
|
finally:
|
|
os.umask(old_umask) # Restore original umask
|
|
return True
|
|
except Exception as e:
|
|
logging.error(f'Failed to save session to {self.session_file}: {str(e)}')
|
|
return False
|
|
|
|
def _load_session(self) -> bool:
|
|
"""Load session data from file.
|
|
|
|
Returns:
|
|
bool: True if file was found and loaded
|
|
"""
|
|
if not os.path.exists(self.session_file):
|
|
logging.info(f'No session file found at {self.session_file}')
|
|
return False
|
|
try:
|
|
# Security check: verify file permissions
|
|
file_stat = os.stat(self.session_file)
|
|
if file_stat.st_mode & 0o077: # Check if other users have permissions
|
|
logging.warning(f'Session file {self.session_file} has insecure permissions')
|
|
return False
|
|
|
|
logging.info(f'restore session from file {self.session_file}')
|
|
with open(self.session_file, 'r') as f:
|
|
session = yaml.safe_load(f)
|
|
|
|
# Extended validation of session data
|
|
required_fields = ['access_token', 'device_id', 'next_batch']
|
|
if not all(field in session for field in required_fields):
|
|
logging.warning('Session file exists but is incomplete')
|
|
return False
|
|
|
|
# Additional integrity check: validate access token format
|
|
if not isinstance(session['access_token'], str) or len(session['access_token']) < 10:
|
|
logging.warning('Session file contains invalid access token')
|
|
return False
|
|
|
|
self.client.access_token = session['access_token']
|
|
self.client.device_id = session['device_id']
|
|
self.client.next_batch = session['next_batch']
|
|
return True
|
|
except Exception as e:
|
|
logging.error(f'Failed to load session from {self.session_file}: {str(e)}')
|
|
return False
|
|
|
|
async def _login(self) -> bool:
|
|
"""Restore login with session file or run full login process.
|
|
|
|
Returns:
|
|
bool: True when bot has a valid login
|
|
"""
|
|
if self._load_session():
|
|
self.session = await self.client.sync(timeout=3000)
|
|
if self.session.transport_response.ok:
|
|
logging.debug('restore session successfully')
|
|
return True
|
|
else:
|
|
logging.warning('no session file was found or file is broken - will attempt full login')
|
|
|
|
logging.debug(f'try login with account {self.username}')
|
|
try:
|
|
self.session = await self.client.login(self.__password)
|
|
if isinstance(self.session, LoginResponse):
|
|
logging.info(f'login successful with account {self.username} on server {self.homeserver}')
|
|
logging.info(self.session)
|
|
self._save_session()
|
|
return True
|
|
else:
|
|
logging.error(f"Login failed for user {self.username}: Unexpected response type {type(self.session)}")
|
|
return False
|
|
except Exception as e:
|
|
logging.error(f"Login attempt failed for user {self.username}: {str(e)}")
|
|
return False
|
|
|
|
def _find_skills(self) -> None:
|
|
"""Search all modules in skill folder and init these."""
|
|
skills_path = pathlib.Path(__file__).parent / "skills"
|
|
for module in skills_path.iterdir():
|
|
if module.is_dir() and (module / "__init__.py").exists():
|
|
importlib.import_module(f"skills.{module.name}")
|
|
|
|
async def _activate_skills(self) -> None:
|
|
"""Activate all skills which defined in config file."""
|
|
for skill in self._config['skills'].keys():
|
|
await self.load_skill(skill_name=skill, config=self._config['skills'][skill])
|
|
|
|
async def load_skill(self, skill_name: str, config: dict = None) -> bool:
|
|
"""Load skill from skill registry.
|
|
|
|
Args:
|
|
skill_name (str): name of skill (same as classname)
|
|
config (dict, optional): additional config. Defaults to None.
|
|
|
|
Returns:
|
|
bool: True is skill is loaded
|
|
"""
|
|
if not skill_name in Skill.registry.keys():
|
|
return False
|
|
logging.info(f'load skill {skill_name} version {Skill.registry[skill_name].__VERSION__}')
|
|
skill = Skill.registry[skill_name](client=self.client, config=config)
|
|
await skill.load()
|
|
self.skills.append(skill)
|
|
return True
|
|
|
|
def unload_skill(self) -> None:
|
|
"""Unload a skill."""
|
|
pass
|
|
|
|
async def send_status(self, message: str) -> bool:
|
|
"""Send a message to a predefined room.
|
|
|
|
Args:
|
|
message (str): message content
|
|
|
|
Returns:
|
|
bool: True if message sent
|
|
"""
|
|
if self.silence:
|
|
return False
|
|
await self.helper.send_message(f"{self._config['general']['status_room']}:{self.domain}", message)
|
|
return True
|
|
|
|
async def run(self) -> None:
|
|
"""Run the bot."""
|
|
if await self._login():
|
|
del self.__password
|
|
await self.send_status(f"{self.botname} Online!")
|
|
await self._activate_skills()
|
|
await self.client.sync_forever(timeout=30000, set_presence='online')
|
|
await self.set_displayname(self._config['general']['botname'])
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
"""evaluates given arguments
|
|
|
|
Returns:
|
|
argsparse.Namespace: Namespace Object with all arguments insert (use: args.long_name_of_argument)
|
|
"""
|
|
argp = argparse.ArgumentParser(description=__doc__)
|
|
# Default args
|
|
argp.add_argument('-v', '--verbose', action='count', default=0,
|
|
help='increase output verbosity (use up to 3 times)')
|
|
argp.add_argument('-c', '--config', default='aoc.conf',
|
|
help='path to config file')
|
|
argp.add_argument('-s', '--session', default='session.tmp',
|
|
help='path to session file')
|
|
|
|
args = argp.parse_args()
|
|
return args
|
|
|
|
async def main():
|
|
args = parse_args()
|
|
if 1 <= args.verbose <= 3:
|
|
logging.getLogger().setLevel(40 - (args.verbose * 10))
|
|
bot = MatrixBot(config_file=args.config, session_file=args.session)
|
|
|
|
bot.silence = True
|
|
await bot.run()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main()) |