README.md

@@ -63,6 +63,8 @@ See `modules/extra_commands.py` and `modules/http_service.py` for an example of
 
 # Todo
 
-- Create module to add auth for the cli connections...
-- Create module for module to module communication...
+- Maybe add ability to search for modules in directories set by a command
 
+- Create module to add auth for the cli connections...
+
+- Add error handling to core_daemon and daemon_cli...

core_daemon.py

@@ -1,77 +1,44 @@
 import os
-import time
 import sys
 import importlib
 import socket
 import json
 import threading
-import logging
-import traceback
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# Add the modules folder to the Python path
+modules_path = os.path.join(os.path.dirname(__file__), 'modules')
+sys.path.append(modules_path)
 
 class ModuleManager:
-    def __init__(self, module_dirs):
-        self.module_dirs = module_dirs
+    def __init__(self):
         self.loaded_modules = {}
         self.extra_commands = {}
-        self._update_sys_path()
-
-    def _update_sys_path(self):
-        for dir in self.module_dirs:
-            full_path = os.path.abspath(dir)
-            if full_path not in sys.path:
-                sys.path.append(full_path)
-                logging.info(f"Added {full_path} to sys.path")
-
-    def add_module_dir(self, new_dir):
-        try:
-            full_path = os.path.abspath(new_dir)
-            if full_path not in self.module_dirs:
-                self.module_dirs.append(full_path)
-                self._update_sys_path()
-                return True, f"Added module directory: {full_path}"
-            return False, f"Module directory already exists: {full_path}"
-        except Exception as e:
-            logging.error(f"Error adding module directory: {str(e)}")
-            return False, f"Error adding module directory: {str(e)}"
 
     def load_module(self, module_name):
-        for dir in self.module_dirs:
-            try:
-                module = importlib.import_module(f'{os.path.basename(dir)}.{module_name}')
-                self.loaded_modules[module_name] = module
-                if hasattr(module, 'initialize'):
-                    module.initialize()
-                if hasattr(module, 'get_commands'):
-                    new_commands = module.get_commands()
-                    self.extra_commands.update(new_commands)
-                logging.info(f"Module '{module_name}' loaded successfully from {dir}.")
-                return True, f"Module '{module_name}' loaded and initialized successfully from {dir}."
-            except ImportError:
-                continue
-            except Exception as e:
-                logging.error(f"Error loading module '{module_name}' from {dir}: {str(e)}")
-                return False, f"Error loading module '{module_name}' from {dir}: {str(e)}"
-        return False, f"Error: Unable to load module '{module_name}' from any of the module directories."
+        try:
+            # Try to import from the modules folder
+            module = importlib.import_module(f'modules.{module_name}')
+            self.loaded_modules[module_name] = module
+            if hasattr(module, 'initialize'):
+                module.initialize()
+            if hasattr(module, 'get_commands'):
+                new_commands = module.get_commands()
+                self.extra_commands.update(new_commands)
+            return True, f"Module '{module_name}' loaded and initialized successfully."
+        except ImportError as e:
+            return False, f"Error: Unable to load module '{module_name}'. {str(e)}"
 
     def unload_module(self, module_name):
         if module_name in self.loaded_modules:
-            try:
-                module = self.loaded_modules[module_name]
-                if hasattr(module, 'shutdown'):
-                    module.shutdown()
-                if hasattr(module, 'get_commands'):
-                    commands_to_remove = module.get_commands().keys()
-                    for cmd in commands_to_remove:
-                        self.extra_commands.pop(cmd, None)
-                del self.loaded_modules[module_name]
-                logging.info(f"Module '{module_name}' unloaded successfully.")
-                return True, f"Module '{module_name}' unloaded and shut down."
-            except Exception as e:
-                logging.error(f"Error unloading module '{module_name}': {str(e)}")
-                return False, f"Error unloading module '{module_name}': {str(e)}"
+            module = self.loaded_modules[module_name]
+            if hasattr(module, 'shutdown'):
+                module.shutdown()
+            if hasattr(module, 'get_commands'):
+                commands_to_remove = module.get_commands().keys()
+                for cmd in commands_to_remove:
+                    self.extra_commands.pop(cmd, None)
+            del self.loaded_modules[module_name]
+            return True, f"Module '{module_name}' unloaded and shut down."
         return False, f"Module '{module_name}' is not loaded."
 
     def list_modules(self):
@@ -82,114 +49,61 @@ class ModuleManager:
 
     def execute_command(self, command, args):
         if command in self.extra_commands:
-            try:
-                result = self.extra_commands[command](args)
-                return True, result
-            except Exception as e:
-                logging.error(f"Error executing command '{command}': {str(e)}")
-                return False, f"Error executing command '{command}': {str(e)}"
+            return True, self.extra_commands[command](args)
         return False, "Command not found"
 
 class CoreDaemon:
-    def __init__(self, host='localhost', port=9999, module_dirs=None):
+    def __init__(self, host='localhost', port=9999):
         self.host = host
         self.port = port
-        self.module_dirs = module_dirs or ['modules']
-        self.module_manager = ModuleManager(self.module_dirs)
+        self.module_manager = ModuleManager()
 
     def start(self):
-        try:
-            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-                s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-                s.bind((self.host, self.port))
-                s.listen()
-                logging.info(f"Core daemon listening on {self.host}:{self.port}")
-                while True:
-                    conn, addr = s.accept()
-                    threading.Thread(target=self.handle_client, args=(conn,)).start()
-        except Exception as e:
-            logging.error(f"Error starting core daemon: {str(e)}")
-            sys.exit(1)
-
-    def shutdown(self):
-        logging.info("Shutting down core daemon...")
-        return {'success': True, 'message': 'Core daemon is shutting down'}
-
-    def _delayed_shutdown(self):
-        time.sleep(1)  # Give time for the response to be sent
-        os._exit(0)
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind((self.host, self.port))
+            s.listen()
+            print(f"Core daemon listening on {self.host}:{self.port}")
+            while True:
+                conn, addr = s.accept()
+                threading.Thread(target=self.handle_client, args=(conn,)).start()
 
     def handle_client(self, conn):
         with conn:
             while True:
-                try:
-                    data = conn.recv(1024)
-                    if not data:
-                        break
-                    command = json.loads(data.decode())
-                    response = self.process_command(command)
-                    conn.sendall(json.dumps(response).encode())
-                except json.JSONDecodeError:
-                    logging.error("Received invalid JSON data")
-                    conn.sendall(json.dumps({'success': False, 'message': 'Invalid command format'}).encode())
-                except Exception as e:
-                    logging.error(f"Error handling client request: {str(e)}")
-                    conn.sendall(json.dumps({'success': False, 'message': 'Internal server error'}).encode())
+                data = conn.recv(1024)
+                if not data:
+                    break
+                command = json.loads(data.decode())
+                response = self.process_command(command)
+                conn.sendall(json.dumps(response).encode())
 
     def upgrade(self):
-        logging.info("Upgrading core daemon...")
-        try:
-            # Here you would typically download or copy the new version
-            # For this example, we'll just restart the script
-            os.execv(sys.executable, ['python'] + sys.argv)
-        except Exception as e:
-            logging.error(f"Error during upgrade: {str(e)}")
-            return False, f"Upgrade failed: {str(e)}"
+        print("Upgrading core daemon...")
+        # Here you would typically download or copy the new version
+        # For this example, we'll just restart the script
+        os.execv(sys.executable, ['python'] + sys.argv)
 
     def process_command(self, command):
         action = command.get('action')
-        try:
-            if action == 'load':
-                return self.module_manager.load_module(command.get('module'))
-            elif action == 'unload':
-                return self.module_manager.unload_module(command.get('module'))
-            elif action == 'list':
-                modules = self.module_manager.list_modules()
-                return True, modules
-            elif action == 'execute':
-                return self.module_manager.execute_command(command.get('command'), command.get('args'))
-            elif action == 'upgrade':
-                return self.upgrade()
-            elif action == 'list_commands':
-                commands = self.module_manager.list_commands()
-                return True, commands
-            elif action == 'add_module_dir':
-                return self.module_manager.add_module_dir(command.get('dir'))
-            elif action == 'shutdown':
-                response = self.shutdown()
-                # After sending the response, exit the program
-                threading.Thread(target=self._delayed_shutdown).start()
-                return response
-            else:
-                return False, "Unknown command"
-        except Exception as e:
-            logging.error(f"Error processing command {action}: {str(e)}")
-            return False, f"Error processing command {action}: {str(e)}"
+        if action == 'load':
+            success, message = self.module_manager.load_module(command.get('module'))
+        elif action == 'unload':
+            success, message = self.module_manager.unload_module(command.get('module'))
+        elif action == 'list':
+            modules = self.module_manager.list_modules()
+            success, message = True, modules
+        elif action == 'execute':
+            success, message = self.module_manager.execute_command(command.get('command'), command.get('args'))
+        elif action == 'upgrade':
+            self.upgrade()
+            success, message = True, "Upgrade initiated"
+        elif action == 'list_commands':
+            commands = self.module_manager.list_commands()
+            success, message = True, commands
+        else:
+            success, message = False, "Unknown command"
+        return {'success': success, 'message': message}
 
 if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Core Daemon with dynamic module loading")
-    parser.add_argument('--host', default='localhost', help='Host to bind the daemon to')
-    parser.add_argument('--port', type=int, default=9999, help='Port to bind the daemon to')
-    parser.add_argument('--module-dirs', nargs='*', default=['modules'], help='Directories to load modules from')
-    
-    args = parser.parse_args()
-
-    try:
-        daemon = CoreDaemon(host=args.host, port=args.port, module_dirs=args.module_dirs)
-        daemon.start()
-    except Exception as e:
-        logging.critical(f"Critical error in core daemon: {str(e)}")
-        logging.debug(traceback.format_exc())
-        sys.exit(1)
+    daemon = CoreDaemon()
+    daemon.start()

daemon_cli.py

@@ -1,7 +1,6 @@
 import cmd
 import socket
 import json
-import time
 import os
 
 class ModuleShell(cmd.Cmd):
@@ -14,41 +13,17 @@ class ModuleShell(cmd.Cmd):
         self.port = port
         self.update_commands()
 
-    def send_command(self, command, retries=3, delay=2):
-        for attempt in range(retries):
-            try:
-                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-                    s.settimeout(5)  # Set a 5-second timeout
-                    s.connect((self.host, self.port))
-                    s.sendall(json.dumps(command).encode())
-                    data = s.recv(1024).decode()
-                    if data:
-                        return json.loads(data)
-                    else:
-                        print("Received empty response. The server might be restarting.")
-                        if command['action'] == 'upgrade':
-                            return {'success': True, 'message': 'Upgrade initiated. The server is restarting.'}
-            except (socket.timeout, ConnectionRefusedError, json.JSONDecodeError) as e:
-                if attempt < retries - 1:
-                    print(f"Connection failed. Retrying in {delay} seconds...")
-                    time.sleep(delay)
-                else:
-                    print(f"Failed to connect after {retries} attempts: {str(e)}")
-                    return {'success': False, 'message': f"Connection error: {str(e)}"}
-        return {'success': False, 'message': "Failed to connect to the server."}
+    def send_command(self, command):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.connect((self.host, self.port))
+            s.sendall(json.dumps(command).encode())
+            return json.loads(s.recv(1024).decode())
 
     def update_commands(self):
-        try:
-            response = self.send_command({'action': 'list_commands'})
-            if isinstance(response, dict) and response.get('success'):
-                self.dynamic_commands = response['message']
-            elif isinstance(response, list):
-                self.dynamic_commands = response
-            else:
-                print("Unexpected response format when updating commands.")
-                self.dynamic_commands = []
-        except Exception as e:
-            print(f"Error updating commands: {str(e)}")
+        response = self.send_command({'action': 'list_commands'})
+        if response['success']:
+            self.dynamic_commands = response['message']
+        else:
             self.dynamic_commands = []
 
     def default(self, line):
@@ -56,13 +31,10 @@ class ModuleShell(cmd.Cmd):
         command = parts[0]
         args = ' '.join(parts[1:])
         response = self.send_command({'action': 'execute', 'command': command, 'args': args})
-        if isinstance(response, dict):
-            if response.get('success'):
-                print(response['message'])
-            else:
-                print(f"Error: {response.get('message', 'Unknown error')}")
+        if response['success']:
+            print(response['message'])
         else:
-            print(f"Unexpected response: {response}")
+            print(f"Error: {response['message']}")
 
     def completenames(self, text, *ignored):
         dotext = 'do_'+text
@@ -72,32 +44,19 @@ class ModuleShell(cmd.Cmd):
     def do_load(self, arg):
         """Load a module: load <module_name>"""
         response = self.send_command({'action': 'load', 'module': arg})
-        if isinstance(response, dict):
-            print(response.get('message', 'Unknown response'))
-        else:
-            print(f"Unexpected response: {response}")
+        print(response['message'])
         self.update_commands()  # Update commands after loading a new module
 
     def do_unload(self, arg):
         """Unload a module: unload <module_name>"""
         response = self.send_command({'action': 'unload', 'module': arg})
-        if isinstance(response, dict):
-            print(response.get('message', 'Unknown response'))
-        else:
-            print(f"Unexpected response: {response}")
+        print(response['message'])
         self.update_commands()  # Update commands after unloading a module
 
     def do_list(self, arg):
         """List all loaded modules"""
         response = self.send_command({'action': 'list'})
-        if isinstance(response, dict) and response.get('success'):
-            modules = response['message']
-        elif isinstance(response, list):
-            modules = response
-        else:
-            print("Unexpected response format when listing modules.")
-            return
-
+        modules = response['message']
         if modules:
             print("Loaded modules:")
             for module in modules:
@@ -118,44 +77,17 @@ class ModuleShell(cmd.Cmd):
         available_modules = [f[:-3] for f in os.listdir(modules_path) if f.endswith('.py') and not f.startswith('__')]
         return [m for m in available_modules if m.startswith(text)]
 
-    def do_add_module_dir(self, arg):
-        """Add a new module directory: add_module_dir <directory_path>"""
-        if not arg:
-            print("Error: Please provide a directory path.")
-            return
-        response = self.send_command({'action': 'add_module_dir', 'dir': arg})
-        if isinstance(response, dict):
-            print(response.get('message', 'Unknown response'))
-        else:
-            print(f"Unexpected response: {response}")
 
     def do_upgrade(self, arg):
         """Upgrade the core daemon"""
         response = self.send_command({'action': 'upgrade'})
-        if isinstance(response, list):
-            try:
-                print(response[0]['message'])
-            except Exception:
-                print('Core daemon response did not make sense.')
-        
-        print("Please reconnect in a few seconds.")
-        return True  # This will exit the command loop
+        print(response['message'])
 
     def do_exit(self, arg):
         """Exit the shell"""
         print("Exiting...")
         return True
 
-    def do_shutdown(self, arg):
-        """Shutdown the core daemon"""
-        response = self.send_command({'action': 'shutdown'})
-        if isinstance(response, dict):
-            print(response.get('message', 'Unknown response'))
-        else:
-            print(f"Unexpected response: {response}")
-        print("Core daemon is shutting down. Exiting CLI.")
-        return True  # This will exit the command loop
-
 def main():
     ModuleShell().cmdloop()