Rework and polish plugins for release

base.py, ob.py, plugins.py, and default.py

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
         }
 
 

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():

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):

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()