diff --git a/src/osintbuddy/elements/base.py b/src/osintbuddy/elements/base.py index 227d37c..0053d09 100644 --- a/src/osintbuddy/elements/base.py +++ b/src/osintbuddy/elements/base.py @@ -19,19 +19,16 @@ class BaseElement(object): """ def __init__(self, **kwargs): self.label: str = '' - self.style: dict = {} self.placeholder: str = '' for key, value in kwargs.items(): - if key == 'label' or key == 'style' or key == 'placeholder': + if key == 'label' or key == 'placeholder': setattr(self, key, value) def _base_blueprint(self): return { 'type': self.node_type, 'label': self.label, - 'placeholder': self.placeholder, - 'style': self.style } diff --git a/src/osintbuddy/ob.py b/src/osintbuddy/ob.py index c6fdacb..9586344 100644 --- a/src/osintbuddy/ob.py +++ b/src/osintbuddy/ob.py @@ -64,9 +64,9 @@ def start_lsp(): def start(): - jedi_language_server = start_lsp() + # jedi_language_server = start_lsp() _print_server_details() - import signal + # import signal import uvicorn uvicorn.run( "osintbuddy.server:app", @@ -79,11 +79,11 @@ def start(): log_level='info' ) - def signal_handler(sig, frame): - jedi_language_server.wait(timeout=1) + # def signal_handler(sig, frame): + # jedi_language_server.wait(timeout=1) - signal.signal(signal.SIGINT, signal_handler) - signal.pause() + # signal.signal(signal.SIGINT, signal_handler) + # signal.pause() def create_plugin_wizard(): diff --git a/src/osintbuddy/plugins.py b/src/osintbuddy/plugins.py index e09f68d..286275f 100755 --- a/src/osintbuddy/plugins.py +++ b/src/osintbuddy/plugins.py @@ -1,4 +1,4 @@ -import os, imp, importlib +import os, imp, importlib, inspect from typing import List, Any, Callable from collections import defaultdict from pydantic import BaseModel, ConfigDict @@ -13,6 +13,20 @@ class OBNode(BaseModel): model_config = OBNodeConfig +def plugin_results_middleman(f): + def return_result(r): + return r + def yield_result(r): + for i in r: + yield i + def decorator(*a, **kwa): + if inspect.isgeneratorfunction(f): + return yield_result(f(*a, **kwa)) + else: + return return_result(f(*a, **kwa)) + return decorator + + class OBAuthorUse(BaseModel): get_driver: Callable[[], None] get_graph: Callable[[], None] @@ -53,7 +67,7 @@ async def get_plugin(cls, plugin_label: str): :return: The plugin class or None if not found. """ for idx, label in enumerate(cls.labels): - if to_snake_case(label) == to_snake_case(plugin_label): + if label == plugin_label or to_snake_case(label) == to_snake_case(plugin_label): return cls.plugins[idx] return None @@ -71,7 +85,7 @@ def get_plug(cls, plugin_label: str): return cls.plugins[idx] return None - def __getitem__(self, i): + def __getitem__(self, i: str): return self.get_plug[i] # https://stackoverflow.com/a/7548190 @@ -168,7 +182,6 @@ class OBPlugin(object, metaclass=OBRegistry): label: str = '' icon: str = 'atom-2' show_label = True - style: dict = {} author = 'Unknown' description = 'No description.' @@ -212,7 +225,6 @@ def blueprint(cls, **kwargs): node['label'] = cls.label node['color'] = cls.color if cls.color else '#145070' node['icon'] = cls.icon - node['style'] = cls.style node['elements'] = [] for element in cls.node: if isinstance(element, list): @@ -275,7 +287,6 @@ def _map_element(transform_map: dict, element: dict): def _map_to_transform_data(cls, node: dict) -> OBNode: transform_map: dict = {} data: dict = node.get('data', {}) - # model_label: str = data.get('label') elements: list[dict] = data.get('elements', []) for element in elements: if isinstance(element, list): diff --git a/src/osintbuddy/templates/default.py b/src/osintbuddy/templates/default.py index 96a342d..e1e455a 100644 --- a/src/osintbuddy/templates/default.py +++ b/src/osintbuddy/templates/default.py @@ -6,7 +6,7 @@ def plugin_source_template(label: str, description: str, author: str) -> str: class {class_name}(ob.Plugin): label = '{label}' - icon = 'atom' # https://tabler-icons.io/ + icon = 'atom-2' # https://tabler-icons.io/ color = '#FFD166' author = '{author}' @@ -16,7 +16,7 @@ class {class_name}(ob.Plugin): TextInput(label='Example', icon='radioactive') ] - @ob.transform(label='To example', icon='atom') + @ob.transform(label='To example', icon='atom-2') async def transform_example(self, node, use): WebsitePlugin = await ob.Registry.get_plugin('website') website_plugin = WebsitePlugin()