chrissewell/pwnagotchi
1
0
Fork 0
mirror of https://github.com/jayofelony/pwnagotchi.git synced 2025-07-01 18:37:27 -04:00
Code Issues Packages Projects Releases Wiki Activity

Update __init__.py

Browse Source
This commit is contained in:
XxKingsxX-Pinu
2024-07-28 17:46:53 +10:00
parent 1ee940c798
commit 0f36a10567
1 changed files with 45 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

182
pwnagotchi/plugins/__init__.py
View File

@ -1,105 +1,17 @@
import os
import queue
import glob
import _thread import _thread
import threading import glob
import importlib, importlib.util import importlib
import importlib.util
import logging import logging
import time import os
import threading
import pwnagotchi.grid
import prctl import prctl
#Idea and base code from NurseJackass
default_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "default") default_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "default")
loaded = {} loaded = {}
database = {} database = {}
locks = {} locks = {}
exitFlag = 0
plugin_event_queues = {}
plugin_thread_workers = {}
class PluginHandler():
def __init__(self, plugin_name):
try:
self._worker_thread = threading.Thread(target=self.doWork, daemon=True, name = "%s.sleeping" % plugin_name)
self._loop_thread = None
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("Starting worker for %s" % plugin_name)
self._worker_thread.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" or event_name == "loop":
# 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
# for future use, use `on_loop`
try:
self._loop_thread = threading.Thread(target=self.doLoop, args = (self, event_name, *args), daemon=True, name = "%s.loop" % (self.plugin_name)).start()
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))
self._worker_thread.start()
logging.info("Worker thread exited 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)
if callback:
callback(*args, **kwargs)
def doWork(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
prctl.set_name("pwnagotchi.%s.%s" % (self.plugin_name, event_name ))
self._worker_thread.name = "%s.%s" % (self.plugin_name, event_name)
logging.debug("")
self.process_event(event_name, *args, **kwargs)
except queue.Empty as e:
self._worker_thread.name = "%s.sleeping"
prctl.set_name("pwnagotchi.%s.sleeping" % (self.plugin_name))
pass
except Exception as e:
logging.exception(repr(e))
def doLoop(self, loopCB, event_name, *args, **kwargs):
global exitFlag
plugin_name = self.plugin_name
prctl.set_name("pwnagotchi.%s" % self.plugin_name)
while not exitFlag and self.keep_going:
try:
self.process_event(event_name, *args, **kwargs)
self.keep_going = False
except Exception as e:
#error in plugin loop kill plugin
self.keep_going = False
logging.exception(repr(e))
def killLoop(self):
self._loop_thread.stop()
class Plugin: class Plugin:
@ -133,19 +45,16 @@ def toggle_plugin(name, enable=True):
global loaded, database global loaded, database
if pwnagotchi.config: if pwnagotchi.config:
if not name in pwnagotchi.config['main']['plugins']: if name not in pwnagotchi.config['main']['plugins']:
pwnagotchi.config['main']['plugins'][name] = dict() pwnagotchi.config['main']['plugins'][name] = dict()
pwnagotchi.config['main']['plugins'][name]['enabled'] = enable pwnagotchi.config['main']['plugins'][name]['enabled'] = enable
save_config(pwnagotchi.config, '/etc/pwnagotchi/config.toml')
if not enable and name in loaded: if not enable and name in loaded:
if getattr(loaded[name], 'on_unload', None): if getattr(loaded[name], 'on_unload', None):
loaded[name].on_unload(view.ROOT) loaded[name].on_unload(view.ROOT)
del loaded[name] 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 return True
if enable and name in database and name not in loaded: if enable and name in database and name not in loaded:
@ -153,44 +62,51 @@ def toggle_plugin(name, enable=True):
if name in loaded and pwnagotchi.config and name in pwnagotchi.config['main']['plugins']: if name in loaded and pwnagotchi.config and name in pwnagotchi.config['main']['plugins']:
loaded[name].options = pwnagotchi.config['main']['plugins'][name] loaded[name].options = pwnagotchi.config['main']['plugins'][name]
one(name, 'loaded') one(name, 'loaded')
time.sleep(3)
if pwnagotchi.config: if pwnagotchi.config:
one(name, 'config_changed', pwnagotchi.config) one(name, 'config_changed', pwnagotchi.config)
one(name, 'ui_setup', view.ROOT) one(name, 'ui_setup', view.ROOT)
one(name, 'ready', view.ROOT._agent) one(name, 'ready', view.ROOT._agent)
if pwnagotchi.config:
save_config(pwnagotchi.config, '/etc/pwnagotchi/config.toml')
return True return True
return False return False
def on(event_name, *args, **kwargs): def on(event_name, *args, **kwargs):
global loaded, plugin_event_queues
cb_name = 'on_%s' % event_name
for plugin_name in loaded.keys(): for plugin_name in loaded.keys():
plugin = loaded[plugin_name]
callback = getattr(plugin, cb_name, None)
if callback is None or not callable(callback): one(plugin_name, event_name, *args, **kwargs)
continue
if plugin_name not in plugin_event_queues:
plugin_event_queues[plugin_name] = PluginHandler(plugin_name)
plugin_event_queues[plugin_name].AddWork(event_name, *args, **kwargs) 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)
def one(plugin_name, event_name, *args, **kwargs): def one(plugin_name, event_name, *args, **kwargs):
global loaded, plugin_event_queues global loaded
if plugin_name in loaded: if plugin_name in loaded:
plugin = loaded[plugin_name] plugin = loaded[plugin_name]
cb_name = 'on_%s' % event_name cb_name = 'on_%s' % event_name
callback = getattr(plugin, cb_name, None) callback = getattr(plugin, cb_name, None)
if callback is not None and callable(callback): if callback is not None and callable(callback):
if plugin_name not in plugin_event_queues: try:
plugin_event_queues[plugin_name] = PluginHandler(plugin_name) 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()
plugin_event_queues[plugin_name].AddWork(event_name, *args, **kwargs) except Exception as e:
logging.error("error while running %s.%s : %s" % (plugin_name, cb_name, e))
logging.error(e, exc_info=True)
def load_from_file(filename): def load_from_file(filename):
@ -199,8 +115,6 @@ def load_from_file(filename):
spec = importlib.util.spec_from_file_location(plugin_name, filename) spec = importlib.util.spec_from_file_location(plugin_name, filename)
instance = importlib.util.module_from_spec(spec) instance = importlib.util.module_from_spec(spec)
spec.loader.exec_module(instance) spec.loader.exec_module(instance)
if plugin_name not in plugin_event_queues:
plugin_event_queues[plugin_name] = PluginHandler(plugin_name)
return plugin_name, instance return plugin_name, instance
@ -221,26 +135,20 @@ def load_from_path(path, enabled=()):
def load(config): def load(config):
try: enabled = [name for name, options in config['main']['plugins'].items() if
enabled = [name for name, options in config['main']['plugins'].items() if 'enabled' in options and options['enabled']]
'enabled' in options and options['enabled']]
# load default plugins # load default plugins
load_from_path(default_path, enabled=enabled) load_from_path(default_path, enabled=enabled)
# load custom ones # load custom ones
custom_path = config['main']['custom_plugins'] if 'custom_plugins' in config['main'] else None custom_path = config['main']['custom_plugins'] if 'custom_plugins' in config['main'] else None
if custom_path is not None: if custom_path is not None:
load_from_path(custom_path, enabled=enabled) load_from_path(custom_path, enabled=enabled)
# propagate options # propagate options
for name, plugin in loaded.items(): for name, plugin in loaded.items():
if name in config['main']['plugins']: plugin.options = config['main']['plugins'][name]
plugin.options = config['main']['plugins'][name]
else:
plugin.options = {}
on('loaded') on('loaded')
on('config_changed', config) on('config_changed', config)
except Exception as e:
logging.exception(repr(e))