395 lines
15 KiB
Python
395 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
|
||
"""
|
||
Funguy Bot Class
|
||
"""
|
||
|
||
import os
|
||
import logging
|
||
import importlib
|
||
import simplematrixbotlib as botlib
|
||
from dotenv import load_dotenv
|
||
import time
|
||
import sys
|
||
import toml
|
||
import socket
|
||
import asyncio
|
||
from collections import defaultdict
|
||
|
||
from plugins.config import FunguyConfig
|
||
|
||
# Rate limiter settings
|
||
RATE_LIMIT_WINDOW = 5.0 # seconds
|
||
MAX_COMMANDS_PER_WINDOW = 3
|
||
|
||
|
||
class FunguyBot:
|
||
"""
|
||
A bot class for managing plugins and handling commands in a Matrix chat environment.
|
||
"""
|
||
|
||
def __init__(self):
|
||
print("[INIT] Starting FunguyBot initialization...")
|
||
|
||
self.PLUGINS_DIR = "plugins"
|
||
self.PLUGINS = {}
|
||
self.config = None
|
||
self.bot = None
|
||
self.disabled_plugins = {}
|
||
|
||
# Rate limiter state: {sender: [(timestamp, room_id), ...]}
|
||
self._rate_limit_buckets = defaultdict(list)
|
||
|
||
load_dotenv() # load once here
|
||
self.setup_logging()
|
||
self.load_plugins()
|
||
self.load_config()
|
||
self.load_disabled_plugins()
|
||
|
||
print("[INIT] FunguyBot initialization complete!")
|
||
|
||
def setup_logging(self):
|
||
log_level = os.getenv("LOG_LEVEL", "INFO").upper()
|
||
level_map = {
|
||
"DEBUG": logging.DEBUG,
|
||
"INFO": logging.INFO,
|
||
"WARNING": logging.WARNING,
|
||
"ERROR": logging.ERROR,
|
||
"CRITICAL": logging.CRITICAL
|
||
}
|
||
level = level_map.get(log_level, logging.INFO)
|
||
|
||
logging.basicConfig(
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||
level=level
|
||
)
|
||
logging.getLogger().setLevel(level)
|
||
logging.getLogger("aiohttp").setLevel(logging.WARNING)
|
||
logging.getLogger("nio").setLevel(logging.WARNING)
|
||
logging.info(f"Logging configured with level: {log_level}")
|
||
|
||
def load_plugins(self):
|
||
for plugin_file in os.listdir(self.PLUGINS_DIR):
|
||
if plugin_file.endswith(".py") and plugin_file != "__init__.py":
|
||
plugin_name = os.path.splitext(plugin_file)[0]
|
||
try:
|
||
module = importlib.import_module(f"{self.PLUGINS_DIR}.{plugin_name}")
|
||
self.PLUGINS[plugin_name] = module
|
||
logging.info(f"Loaded plugin: {plugin_name}")
|
||
except Exception as e:
|
||
logging.error(f"Error loading plugin {plugin_name}: {e}")
|
||
|
||
def setup_plugins(self):
|
||
"""Call setup(bot) on any plugin that defines it, after self.bot exists."""
|
||
for plugin_name, plugin_module in self.PLUGINS.items():
|
||
if hasattr(plugin_module, "setup") and callable(plugin_module.setup):
|
||
try:
|
||
plugin_module.setup(self.bot)
|
||
logging.info(f"Setup called for plugin: {plugin_name}")
|
||
except Exception as e:
|
||
logging.error(f"Error during setup of plugin {plugin_name}: {e}")
|
||
|
||
def reload_plugins(self):
|
||
self.PLUGINS.clear()
|
||
for plugin_name in list(sys.modules.keys()):
|
||
if plugin_name.startswith(self.PLUGINS_DIR + "."):
|
||
del sys.modules[plugin_name]
|
||
self.load_plugins()
|
||
if self.bot is not None:
|
||
self.setup_plugins()
|
||
|
||
def load_config(self):
|
||
self.config = FunguyConfig()
|
||
logging.info("Configuration loaded")
|
||
|
||
def load_disabled_plugins(self):
|
||
if os.path.exists('funguy.conf'):
|
||
with open('funguy.conf', 'r') as f:
|
||
config_data = toml.load(f)
|
||
self.disabled_plugins = config_data.get('plugins', {}).get('disabled', {})
|
||
|
||
def save_disabled_plugins(self):
|
||
existing_config = {}
|
||
if os.path.exists('funguy.conf'):
|
||
with open('funguy.conf', 'r') as f:
|
||
existing_config = toml.load(f)
|
||
existing_config['plugins'] = {'disabled': self.disabled_plugins}
|
||
with open('funguy.conf', 'w') as f:
|
||
toml.dump(existing_config, f)
|
||
|
||
def _check_rate_limit(self, sender: str) -> bool:
|
||
"""Return True if the sender is allowed to proceed.
|
||
Admin is always allowed."""
|
||
# Admin bypass
|
||
if sender == self.config.admin_user:
|
||
return True
|
||
|
||
now = time.monotonic()
|
||
bucket = self._rate_limit_buckets[sender]
|
||
# Prune old entries
|
||
bucket = [t for t in bucket if now - t < RATE_LIMIT_WINDOW]
|
||
self._rate_limit_buckets[sender] = bucket
|
||
|
||
if len(bucket) >= MAX_COMMANDS_PER_WINDOW:
|
||
logging.debug("Rate limit hit for %s", sender)
|
||
return False
|
||
|
||
bucket.append(now)
|
||
return True
|
||
|
||
# ------------------------------------------------------------------
|
||
# New: load/unload a single plugin at runtime
|
||
# ------------------------------------------------------------------
|
||
async def load_plugin(self, plugin_name: str) -> bool:
|
||
"""Dynamically load a plugin module, add to PLUGINS, and call its setup()."""
|
||
if plugin_name in self.PLUGINS:
|
||
logging.info(f"Plugin '{plugin_name}' is already loaded.")
|
||
return False
|
||
try:
|
||
module = importlib.import_module(f"{self.PLUGINS_DIR}.{plugin_name}")
|
||
self.PLUGINS[plugin_name] = module
|
||
logging.info(f"Loaded plugin: {plugin_name}")
|
||
# Call setup if the bot is already running
|
||
if self.bot is not None and hasattr(module, "setup") and callable(module.setup):
|
||
module.setup(self.bot)
|
||
logging.info(f"Setup called for newly loaded plugin: {plugin_name}")
|
||
return True
|
||
except Exception as e:
|
||
logging.error(f"Error loading plugin {plugin_name}: {e}")
|
||
return False
|
||
|
||
async def unload_plugin(self, plugin_name: str) -> bool:
|
||
"""Remove a plugin from PLUGINS and unload its module."""
|
||
if plugin_name not in self.PLUGINS:
|
||
logging.info(f"Plugin '{plugin_name}' is not loaded.")
|
||
return False
|
||
try:
|
||
del self.PLUGINS[plugin_name]
|
||
module_path = f"{self.PLUGINS_DIR}.{plugin_name}"
|
||
if module_path in sys.modules:
|
||
del sys.modules[module_path]
|
||
logging.info(f"Unloaded plugin: {plugin_name}")
|
||
return True
|
||
except Exception as e:
|
||
logging.error(f"Error unloading plugin {plugin_name}: {e}")
|
||
return False
|
||
|
||
# ------------------------------------------------------------------
|
||
# New: restart the bot process
|
||
# ------------------------------------------------------------------
|
||
async def restart_bot(self, room_id):
|
||
await self.bot.api.send_text_message(room_id, "🔄 Restarting bot...")
|
||
await asyncio.sleep(1)
|
||
logging.info("Restart command received – exiting.")
|
||
sys.exit(0)
|
||
|
||
async def handle_commands(self, room, message):
|
||
match = botlib.MessageMatch(room, message, self.bot, self.config.prefix)
|
||
|
||
sender = str(message.sender)
|
||
is_admin = (sender == self.config.admin_user)
|
||
|
||
# Rate limit check (applies to all non‑admin commands)
|
||
if not self._check_rate_limit(sender):
|
||
await self.bot.api.send_text_message(
|
||
room.room_id,
|
||
"⛔ You're sending commands too quickly. Please wait a few seconds."
|
||
)
|
||
return
|
||
|
||
# Admin commands
|
||
if match.is_not_from_this_bot() and match.prefix() and match.command("reload"):
|
||
if is_admin:
|
||
self.reload_plugins()
|
||
await self.bot.api.send_text_message(room.room_id, "All plugins reloaded successfully.")
|
||
else:
|
||
await self.bot.api.send_text_message(room.room_id, "You are not authorized to reload plugins.")
|
||
return
|
||
|
||
if match.is_not_from_this_bot() and match.prefix() and match.command("load"):
|
||
if not is_admin:
|
||
await self.bot.api.send_text_message(room.room_id, "⛔ Admin only.")
|
||
return
|
||
args = match.args()
|
||
if len(args) != 1:
|
||
await self.bot.api.send_text_message(room.room_id, "Usage: !load <plugin>")
|
||
return
|
||
success = await self.load_plugin(args[0])
|
||
msg = f"✅ Plugin '{args[0]}' loaded." if success else f"❌ Could not load '{args[0]}'. See logs for details."
|
||
await self.bot.api.send_text_message(room.room_id, msg)
|
||
return
|
||
|
||
if match.is_not_from_this_bot() and match.prefix() and match.command("unload"):
|
||
if not is_admin:
|
||
await self.bot.api.send_text_message(room.room_id, "⛔ Admin only.")
|
||
return
|
||
args = match.args()
|
||
if len(args) != 1:
|
||
await self.bot.api.send_text_message(room.room_id, "Usage: !unload <plugin>")
|
||
return
|
||
success = await self.unload_plugin(args[0])
|
||
msg = f"✅ Plugin '{args[0]}' unloaded." if success else f"❌ Could not unload '{args[0]}'. See logs for details."
|
||
await self.bot.api.send_text_message(room.room_id, msg)
|
||
return
|
||
|
||
if match.is_not_from_this_bot() and match.prefix() and match.command("disable"):
|
||
if not is_admin:
|
||
await self.bot.api.send_text_message(room.room_id, "⛔ Admin only.")
|
||
return
|
||
args = match.args()
|
||
if len(args) != 1:
|
||
await self.bot.api.send_text_message(room.room_id, "Usage: !disable <plugin>")
|
||
return
|
||
plugin_name = args[0]
|
||
room_id = room.room_id
|
||
await self.disable_plugin(room_id, plugin_name)
|
||
await self.bot.api.send_text_message(room.room_id, f"🚫 Plugin '{plugin_name}' disabled in this room.")
|
||
return
|
||
|
||
if match.is_not_from_this_bot() and match.prefix() and match.command("enable"):
|
||
if not is_admin:
|
||
await self.bot.api.send_text_message(room.room_id, "⛔ Admin only.")
|
||
return
|
||
args = match.args()
|
||
if len(args) != 1:
|
||
await self.bot.api.send_text_message(room.room_id, "Usage: !enable <plugin>")
|
||
return
|
||
plugin_name = args[0]
|
||
room_id = room.room_id
|
||
await self.enable_plugin(room_id, plugin_name)
|
||
await self.bot.api.send_text_message(room.room_id, f"✅ Plugin '{plugin_name}' enabled in this room.")
|
||
return
|
||
|
||
if match.is_not_from_this_bot() and match.prefix() and match.command("restart"):
|
||
if not is_admin:
|
||
await self.bot.api.send_text_message(room.room_id, "⛔ Admin only.")
|
||
return
|
||
await self.restart_bot(room.room_id)
|
||
return
|
||
|
||
if match.is_not_from_this_bot() and match.prefix() and match.command("rehash"):
|
||
if not is_admin:
|
||
await self.bot.api.send_text_message(room.room_id, "You are not authorized to reload config.")
|
||
return
|
||
self.load_config()
|
||
self.load_disabled_plugins()
|
||
await self.bot.api.send_text_message(room.room_id, "🔄 Configuration rehashed.")
|
||
return
|
||
|
||
# Dispatch to active plugins
|
||
for plugin_name, plugin_module in self.PLUGINS.items():
|
||
if plugin_name not in self.disabled_plugins.get(room.room_id, []):
|
||
try:
|
||
await plugin_module.handle_command(room, message, self.bot, self.config.prefix, self.config)
|
||
except Exception as e:
|
||
logging.error(f"Error in plugin {plugin_name}: {e}", exc_info=True)
|
||
|
||
async def disable_plugin(self, room_id, plugin_name):
|
||
if room_id not in self.disabled_plugins:
|
||
self.disabled_plugins[room_id] = []
|
||
if plugin_name not in self.disabled_plugins[room_id]:
|
||
self.disabled_plugins[room_id].append(plugin_name)
|
||
self.save_disabled_plugins()
|
||
|
||
async def enable_plugin(self, room_id, plugin_name):
|
||
if room_id in self.disabled_plugins and plugin_name in self.disabled_plugins[room_id]:
|
||
self.disabled_plugins[room_id].remove(plugin_name)
|
||
self.save_disabled_plugins()
|
||
|
||
def test_connectivity(self, hostname, port=443):
|
||
logging.info(f"Testing connectivity to {hostname}:{port}...")
|
||
try:
|
||
ip_address = socket.gethostbyname(hostname)
|
||
logging.info(f"✓ DNS resolution successful: {hostname} -> {ip_address}")
|
||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
sock.settimeout(10)
|
||
result = sock.connect_ex((hostname, port))
|
||
sock.close()
|
||
if result == 0:
|
||
logging.info(f"✓ Socket connection successful to {hostname}:{port}")
|
||
return True
|
||
else:
|
||
logging.error(f"✗ Socket connection failed to {hostname}:{port} (error code: {result})")
|
||
return False
|
||
except socket.gaierror as e:
|
||
logging.error(f"✗ DNS resolution failed for {hostname}: {e}")
|
||
return False
|
||
except Exception as e:
|
||
logging.error(f"✗ Connectivity test failed: {e}")
|
||
return False
|
||
|
||
def run(self):
|
||
print("\n" + "="*60)
|
||
print("FUNGUY BOT - STARTING")
|
||
print("="*60 + "\n")
|
||
|
||
MATRIX_URL = os.getenv("MATRIX_URL")
|
||
MATRIX_USER = os.getenv("MATRIX_USER")
|
||
MATRIX_PASS = os.getenv("MATRIX_PASS")
|
||
|
||
if not MATRIX_URL or not MATRIX_USER or not MATRIX_PASS:
|
||
logging.error("Missing MATRIX_URL / MATRIX_USER / MATRIX_PASS in .env")
|
||
return
|
||
|
||
logging.info(f"Matrix URL: {MATRIX_URL}")
|
||
logging.info(f"Matrix User: {MATRIX_USER}")
|
||
|
||
hostname = MATRIX_URL.replace("https://", "").replace("http://", "").split("/")[0]
|
||
|
||
logging.info("="*40)
|
||
logging.info("RUNNING NETWORK DIAGNOSTICS")
|
||
logging.info("="*40)
|
||
|
||
if not self.test_connectivity(hostname, 443):
|
||
logging.error("Connectivity test failed. See above messages.")
|
||
return
|
||
|
||
logging.info("="*40)
|
||
logging.info("ATTEMPTING MATRIX CONNECTION")
|
||
logging.info("="*40)
|
||
|
||
try:
|
||
creds = botlib.Creds(MATRIX_URL, MATRIX_USER, MATRIX_PASS)
|
||
self.bot = botlib.Bot(creds, self.config)
|
||
|
||
self.setup_plugins()
|
||
self.bot.plugins = self.PLUGINS
|
||
|
||
@self.bot.listener.on_message_event
|
||
async def wrapper_handle_commands(room, message):
|
||
await self.handle_commands(room, message)
|
||
|
||
self.bot.run()
|
||
except Exception as e:
|
||
logging.error(f"Fatal error during bot startup: {e}", exc_info=True)
|
||
raise
|
||
|
||
def stop(self):
|
||
"""Cleanup resources before shutdown."""
|
||
if hasattr(self, 'bot') and self.bot is not None:
|
||
# try to stop any schedulers if needed
|
||
pass
|
||
logging.info("Bot stopped.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("\n" + "="*60)
|
||
print("FUNGUY BOT LAUNCHER")
|
||
print("="*60)
|
||
|
||
bot = None
|
||
try:
|
||
bot = FunguyBot()
|
||
bot.run()
|
||
except KeyboardInterrupt:
|
||
print("\n[!] Bot stopped by user")
|
||
if bot:
|
||
bot.stop()
|
||
sys.exit(0)
|
||
except Exception as e:
|
||
print(f"\n[!] Fatal error: {e}")
|
||
logging.error(f"Unhandled exception: {e}", exc_info=True)
|
||
if bot:
|
||
bot.stop()
|
||
sys.exit(1)
|