|
@@ -1,61 +1,92 @@
|
|
|
import importlib
|
|
|
+import logging
|
|
|
+import time
|
|
|
+from threading import Thread
|
|
|
|
|
|
-import db
|
|
|
-from config import PLUGINS, ALIASES, RESPONSES, BOTS_LOOKUP
|
|
|
-from command_system import RollbotResponse, RollbotFailure
|
|
|
+from command_system import RollbotResponse, RollbotFailure, as_plugin
|
|
|
+
|
|
|
+
|
|
|
+def lift_response(call, response):
|
|
|
+ @as_plugin(call)
|
|
|
+ def response_func(db, msg):
|
|
|
+ return RollbotResponse(msg, txt=response)
|
|
|
+ return response_func
|
|
|
|
|
|
|
|
|
class Rollbot:
|
|
|
- def __init__(self, logger):
|
|
|
+ def __init__(self, logger=logging.getLogger(__name__), plugins={}, aliases={}, responses={}, lookup={}, sleep_time=0.0, session_factory=None, callback=None):
|
|
|
self.logger = logger
|
|
|
+ self.session_factory = session_factory or (lambda: None)
|
|
|
+ self.callback = callback or (lambda txt, gid: self.logger.info(f"Responding to {gid} with {txt}"))
|
|
|
self.commands = {}
|
|
|
+ self.to_start = set()
|
|
|
+ self.to_stop = set()
|
|
|
+
|
|
|
self.logger.info("Loading command plugins")
|
|
|
- for module, classes in PLUGINS.items():
|
|
|
+ for module, classes in plugins.items():
|
|
|
plugin_module = importlib.import_module("plugins." + module)
|
|
|
for class_name in classes:
|
|
|
+ self.logger.info(class_name)
|
|
|
plugin_class = getattr(plugin_module, class_name)
|
|
|
- plugin_instance = plugin_class(logger)
|
|
|
+ plugin_instance = plugin_class(logger=logger)
|
|
|
if plugin_instance.command in self.commands:
|
|
|
self.logger.error(f"Duplicate command word '{plugin_instance.command}'")
|
|
|
raise ValueError(f"Duplicate command word '{plugin_instance.command}'")
|
|
|
self.commands[plugin_instance.command] = plugin_instance
|
|
|
+ if "on_start" in plugin_class.__dict__:
|
|
|
+ self.to_start.add(plugin_instance)
|
|
|
+ if "on_shutdown" in plugin_class.__dict__:
|
|
|
+ self.to_stop.add(plugin_instance)
|
|
|
self.logger.info(f"Finished loading plugins, {len(self.commands)} commands found")
|
|
|
- self.logger.info("Initializing database tables")
|
|
|
- db.init_db()
|
|
|
- self.logger.info("Finished initializing database")
|
|
|
+
|
|
|
+ self.logger.info("Loading simple responses")
|
|
|
+ for cmd, response in responses.items():
|
|
|
+ if cmd in self.commands:
|
|
|
+ self.logger.error(f"Duplicate command word '{cmd}'")
|
|
|
+ raise ValueError(f"Duplicate command word '{cmd}'")
|
|
|
+ self.commands[cmd] = lift_response(cmd, response)(logger=logger)
|
|
|
+ self.logger.info(f"Finished loading simple responses, {len(self.commands)} total commands available")
|
|
|
+
|
|
|
+ self.logger.info("Loading aliases")
|
|
|
+ for alias, cmd in aliases.items():
|
|
|
+ if cmd not in self.commands:
|
|
|
+ self.logger.error(f"Missing aliased command word '{cmd}'")
|
|
|
+ raise ValueError(f"Missing aliased command word '{cmd}'")
|
|
|
+ if alias in self.commands:
|
|
|
+ self.logger.error(f"Duplicate command word '{alias}'")
|
|
|
+ raise ValueError(f"Duplicate command word '{alias}'")
|
|
|
+ self.commands[alias] = self.commands[cmd]
|
|
|
+ self.logger.info(f"Finished loading aliases, {len(self.commands)} total commands + aliases available")
|
|
|
+
|
|
|
+ self.bot_lookup = lookup
|
|
|
+ self.sleep_time = sleep_time
|
|
|
|
|
|
def start_plugins(self):
|
|
|
- self.logger.info("Starting all plugins")
|
|
|
- with db.session_scope() as session:
|
|
|
- for plugin_instance in self.commands.values():
|
|
|
- plugin_instance.on_start(session)
|
|
|
+ self.logger.info("Starting plugins")
|
|
|
+ with self.session_factory() as session:
|
|
|
+ for cmd in self.to_start:
|
|
|
+ cmd.on_start(session)
|
|
|
self.logger.info("Finished starting plugins")
|
|
|
|
|
|
def shutdown_plugins(self):
|
|
|
- self.logger.info("Shutting down all plugins")
|
|
|
- with db.session_scope() as session:
|
|
|
- for plugin_instance in self.commands.values():
|
|
|
- plugin_instance.on_shutdown(session)
|
|
|
+ self.logger.info("Shutting down plugins")
|
|
|
+ with self.session_factory() as session:
|
|
|
+ for cmd in self.to_stop:
|
|
|
+ cmd.on_shutdown(session)
|
|
|
self.logger.info("Finished shutting down plugins")
|
|
|
|
|
|
def run_command(self, message):
|
|
|
if not message.is_command:
|
|
|
self.logger.warn(f"Tried to run non-command message {message.message_id}")
|
|
|
- return
|
|
|
-
|
|
|
- # if this command is aliased, resolve that first, otherwise use the literal command
|
|
|
- cmd = ALIASES.get(message.command, message.command)
|
|
|
+ return RollbotResponse(msg, failure=RollbotFailure.INTERNAL_ERROR)
|
|
|
|
|
|
- if cmd in RESPONSES:
|
|
|
- return RollbotResponse(message, txt=RESPONSES[cmd])
|
|
|
-
|
|
|
- plugin = self.commands.get(cmd, None)
|
|
|
+ plugin = self.commands.get(message.command, None)
|
|
|
|
|
|
if plugin is None:
|
|
|
- self.logger.warn(f"Message {message.message_id} had a command {message.command} (resolved to {cmd}) that could not be run.")
|
|
|
+ self.logger.warn(f"Message {message.message_id} had a command {message.command} that could not be run.")
|
|
|
return RollbotResponse(message, failure=RollbotFailure.INVALID_COMMAND)
|
|
|
|
|
|
- with db.session_scope() as session:
|
|
|
+ with self.session_factory() as session:
|
|
|
response = plugin.on_command(session, message)
|
|
|
|
|
|
if not response.is_success:
|
|
@@ -63,3 +94,44 @@ class Rollbot:
|
|
|
self.logger.warn(response.info)
|
|
|
|
|
|
return response
|
|
|
+
|
|
|
+ def handle_command(self, msg):
|
|
|
+ if msg.group_id not in self.bot_lookup:
|
|
|
+ self.logger.warning(f"Received message from unknown group ID {msg.group_id}")
|
|
|
+ return
|
|
|
+
|
|
|
+ self.logger.info(f"Handling message {msg.message_id}")
|
|
|
+ t = time.time()
|
|
|
+ try:
|
|
|
+ response = self.run_command(msg)
|
|
|
+ except Exception as e:
|
|
|
+ self.logger.exception(f"Exception during command execution for message {msg.message_id}")
|
|
|
+ response = RollbotResponse(msg, failure=RollbotFailure.INTERNAL_ERROR)
|
|
|
+
|
|
|
+ if not response.respond:
|
|
|
+ self.logger.info(f"Skipping response to message {msg.message_id}")
|
|
|
+ return
|
|
|
+
|
|
|
+ self.logger.info(f"Responding to message {msg.message_id}")
|
|
|
+
|
|
|
+ sleep = self.sleep_time - time.time() + t
|
|
|
+ if sleep > 0:
|
|
|
+ self.logger.info(f"Sleeping for {sleep:.3f}s before responding")
|
|
|
+ time.sleep(sleep)
|
|
|
+
|
|
|
+ bot_id = self.bot_lookup[msg.group_id]
|
|
|
+ if response.is_success:
|
|
|
+ if response.txt is not None:
|
|
|
+ self.callback(response.txt, bot_id)
|
|
|
+ if response.img is not None:
|
|
|
+ self.callback(response.img, bot_id)
|
|
|
+ else:
|
|
|
+ self.callback(response.failure_msg, bot_id)
|
|
|
+ self.logger.warning(f"Failed command response: {response}")
|
|
|
+
|
|
|
+ t = time.time() - t
|
|
|
+ self.logger.info(f"Exiting command thread for {msg.message_id} after {t:.3f}s")
|
|
|
+
|
|
|
+ def handle_command_threaded(self, message):
|
|
|
+ t = Thread(target=lambda: self.handle_command(message))
|
|
|
+ t.start()
|