diff --git a/pwnagotchi/plugins/__init__.py b/pwnagotchi/plugins/__init__.py index 09074d93..76798260 100644 --- a/pwnagotchi/plugins/__init__.py +++ b/pwnagotchi/plugins/__init__.py @@ -1,18 +1,103 @@ -import _thread -import glob -import importlib -import importlib.util -import logging import os +import queue +import glob +import _thread import threading -import pwnagotchi.grid +import importlib, importlib.util +import logging +import time import prctl + default_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "default") loaded = {} database = {} locks = {} +exitFlag = 0 +plugin_event_queues = {} +plugin_thread_workers = {} +def dummy_callback(): + pass + +# callback to run "on_load" in a separate thread for old plugins +# that use on_load like main() and don't return from on_load +# until they are unloading +def run_once(pqueue, event_name, *args, **kwargs): + try: + prctl.set_name("R1_%s_%s" % (pqueue.plugin_name, event_name)) + pqueue.process_event(event_name, *args, *kwargs) + logging.debug("Thread for %s %s exiting" % (pqueue.plugin_name, event_name)) + except Exception as e: + logging.exception("Thread for %s, %s, %s, %s" % (pqueue.plugin_name, event_name, repr(args), repr(kwargs))) + +class PluginEventQueue(threading.Thread): + def __init__(self, plugin_name): + try: + self._worker_thread = threading.Thread.__init__(self, daemon=True) + self.plugin_name = plugin_name + self.work_queue = queue.Queue() + self.queue_lock = threading.Lock() + self.load_handler = None + self.keep_going = True + logging.debug("PLUGIN EVENT QUEUE FOR %s starting %s" % (plugin_name, repr(self.load_handler))) + self.start() + except Exception as e: + logging.exception(e) + + def __del__(self): + self.keep_going = False + self._worker_thread.join() + if self.load_handler: + self.load_handler.join() + + def AddWork(self, event_name, *args, **kwargs): + if event_name == "loaded": + # spawn separate thread, because many plugins use on_load as a "main" loop + # this way on_load can continue if it needs, while other events get processed + try: + cb_name = 'on_%s' % event_name + callback = getattr(loaded[self.plugin_name], cb_name, None) + if callback: + self.load_handler = threading.Thread(target=run_once, + args=(self, event_name, *args), + kwargs=kwargs, + daemon=True) + self.load_handler.start() + else: + self.load_handler = None + except Exception as e: + logging.exception(e) + else: + self.work_queue.put([event_name, args, kwargs]) + + def run(self): + logging.debug("Worker thread starting for %s"%(self.plugin_name)) + prctl.set_name("PLG %s" % self.plugin_name) + self.process_events() + logging.info("Worker thread exiting for %s"%(self.plugin_name)) + + def process_event(self, event_name, *args, **kwargs): + cb_name = 'on_%s' % event_name + callback = getattr(loaded[self.plugin_name], cb_name, None) + logging.debug("%s.%s: %s" % (self.plugin_name, event_name, repr(args))) + if callback: + callback(*args, **kwargs) + + def process_events(self): + global exitFlag + plugin_name = self.plugin_name + work_queue = self.work_queue + + while not exitFlag and self.keep_going: + try: + data = work_queue.get(timeout=2) + (event_name, args, kwargs) = data + self.process_event(event_name, *args, **kwargs) + except queue.Empty as e: + pass + except Exception as e: + logging.exception(repr(e)) class Plugin: @classmethod @@ -45,16 +130,19 @@ def toggle_plugin(name, enable=True): global loaded, database if pwnagotchi.config: - if name not in pwnagotchi.config['main']['plugins']: + if not name in pwnagotchi.config['main']['plugins']: pwnagotchi.config['main']['plugins'][name] = dict() pwnagotchi.config['main']['plugins'][name]['enabled'] = enable - save_config(pwnagotchi.config, '/etc/pwnagotchi/config.toml') if not enable and name in loaded: if getattr(loaded[name], 'on_unload', None): loaded[name].on_unload(view.ROOT) del loaded[name] - + if name in plugin_event_queues: + plugin_event_queues[name].keep_going = False + del plugin_event_queues[name] + if pwnagotchi.config: + save_config(pwnagotchi.config, '/etc/pwnagotchi/config.toml') return True if enable and name in database and name not in loaded: @@ -62,51 +150,45 @@ def toggle_plugin(name, enable=True): if name in loaded and pwnagotchi.config and name in pwnagotchi.config['main']['plugins']: loaded[name].options = pwnagotchi.config['main']['plugins'][name] one(name, 'loaded') + time.sleep(3) if pwnagotchi.config: one(name, 'config_changed', pwnagotchi.config) one(name, 'ui_setup', view.ROOT) one(name, 'ready', view.ROOT._agent) + if pwnagotchi.config: + save_config(pwnagotchi.config, '/etc/pwnagotchi/config.toml') return True return False def on(event_name, *args, **kwargs): + global loaded, plugin_event_queues + cb_name = 'on_%s' % event_name for plugin_name in loaded.keys(): - - one(plugin_name, event_name, *args, **kwargs) + plugin = loaded[plugin_name] + callback = getattr(plugin, cb_name, None) + if callback is None or not callable(callback): + continue -def locked_cb(lock_name, cb, *args, **kwargs): - global locks - - if lock_name not in locks: - locks[lock_name] = threading.Lock() - - with locks[lock_name]: - # Setting the thread name using prctl - plugin_name, plugin_cb = lock_name.split("::") - prctl.set_name(f"{plugin_name}.{plugin_cb}") - cb(*args, **kwargs) + if plugin_name not in plugin_event_queues: + plugin_event_queues[plugin_name] = PluginEventQueue(plugin_name) + plugin_event_queues[plugin_name].AddWork(event_name, *args, **kwargs) + logging.debug("%s %s" % (plugin_name, cb_name)) def one(plugin_name, event_name, *args, **kwargs): - global loaded - + global loaded, plugin_event_queues if plugin_name in loaded: plugin = loaded[plugin_name] cb_name = 'on_%s' % event_name callback = getattr(plugin, cb_name, None) if callback is not None and callable(callback): - try: - lock_name = "%s::%s" % (plugin_name, cb_name) - thread_name = f'{plugin_name}.{cb_name}' - thread = threading.Thread(target=locked_cb, args=(lock_name, callback, *args, *kwargs), name=thread_name, daemon=True) - thread.start() + if plugin_name not in plugin_event_queues: + plugin_event_queues[plugin_name] = PluginEventQueue(plugin_name) - except Exception as e: - logging.error("error while running %s.%s : %s" % (plugin_name, cb_name, e)) - logging.error(e, exc_info=True) + plugin_event_queues[plugin_name].AddWork(event_name, *args, **kwargs) def load_from_file(filename): @@ -115,6 +197,8 @@ def load_from_file(filename): spec = importlib.util.spec_from_file_location(plugin_name, filename) instance = importlib.util.module_from_spec(spec) spec.loader.exec_module(instance) + if plugin_name not in plugin_event_queues: + plugin_event_queues[plugin_name] = PluginEventQueue(plugin_name) return plugin_name, instance @@ -135,20 +219,26 @@ def load_from_path(path, enabled=()): def load(config): - enabled = [name for name, options in config['main']['plugins'].items() if - 'enabled' in options and options['enabled']] + try: + enabled = [name for name, options in config['main']['plugins'].items() if + 'enabled' in options and options['enabled']] - # load default plugins - load_from_path(default_path, enabled=enabled) + # load default plugins + load_from_path(default_path, enabled=enabled) - # load custom ones - custom_path = config['main']['custom_plugins'] if 'custom_plugins' in config['main'] else None - if custom_path is not None: - load_from_path(custom_path, enabled=enabled) + # load custom ones + custom_path = config['main']['custom_plugins'] if 'custom_plugins' in config['main'] else None + if custom_path is not None: + load_from_path(custom_path, enabled=enabled) - # propagate options - for name, plugin in loaded.items(): - plugin.options = config['main']['plugins'][name] + # propagate options + for name, plugin in loaded.items(): + if name in config['main']['plugins']: + plugin.options = config['main']['plugins'][name] + else: + plugin.options = {} - on('loaded') - on('config_changed', config) \ No newline at end of file + on('loaded') + on('config_changed', config) + except Exception as e: + logging.exception(repr(e)) diff --git a/pwnagotchi/ui/hw/gamepi15.py b/pwnagotchi/ui/hw/gamepi15.py index 1a4b54a5..6ef1ad23 100644 --- a/pwnagotchi/ui/hw/gamepi15.py +++ b/pwnagotchi/ui/hw/gamepi15.py @@ -32,7 +32,7 @@ class GamePi15(DisplayImpl): return self._layout def initialize(self): - logging.info("initializing gamepi20 display") + logging.info("initializing gamepi15 display") from pwnagotchi.ui.hw.libs.pimoroni.displayhatmini.ST7789 import ST7789 self._display = ST7789(0, 0, 25, 24, 27, 240, 240, 270, True, 60 * 1000 * 1000, 0, 0) diff --git a/pwnagotchi/ui/hw/spotpear154lcd.py b/pwnagotchi/ui/hw/spotpear154lcd.py index ff382877..5ae5e165 100644 --- a/pwnagotchi/ui/hw/spotpear154lcd.py +++ b/pwnagotchi/ui/hw/spotpear154lcd.py @@ -32,7 +32,7 @@ class Spotpear154lcd(DisplayImpl): return self._layout def initialize(self): - logging.info("initializing gamepi15 display") + logging.info("initializing spotpear154lcd display") from pwnagotchi.ui.hw.libs.pimoroni.displayhatmini.ST7789 import ST7789 self._display = ST7789(0, 0, 22, 24, 27, 240, 240, 0, True, 60 * 1000 * 1000, 0, 0)