upgrade to avoid thread overflow

this is creating a thread per plugin. This is avoiding pressure on the system.

The fix was found by Nursejackass
https://github.com/Sniffleupagus/pwnagotchi-snflpgs/blob/snflpgs/pwnagotchi/plugins/__init__.py
This commit is contained in:
V0r-T3x
2024-12-01 20:37:12 -05:00
parent 83e9ee6318
commit 260e501628

View File

@ -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)
on('loaded')
on('config_changed', config)
except Exception as e:
logging.exception(repr(e))