Update README todo list.

README.md

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

core_daemon.py

@@ -1,44 +1,77 @@
 import os
+import time
 import sys
 import importlib
 import socket
 import json
 import threading
+import logging
+import traceback
 
-# Add the modules folder to the Python path
-modules_path = os.path.join(os.path.dirname(__file__), 'modules')
-sys.path.append(modules_path)
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 class ModuleManager:
-    def __init__(self):
+    def __init__(self, module_dirs):
+        self.module_dirs = module_dirs
         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):
-        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)}"
+        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."
 
     def unload_module(self, module_name):
         if module_name in self.loaded_modules:
-            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."
+            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)}"
         return False, f"Module '{module_name}' is not loaded."
 
     def list_modules(self):
@@ -49,61 +82,114 @@ class ModuleManager:
 
     def execute_command(self, command, args):
         if command in self.extra_commands:
-            return True, self.extra_commands[command](args)
+            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 False, "Command not found"
 
 class CoreDaemon:
-    def __init__(self, host='localhost', port=9999):
+    def __init__(self, host='localhost', port=9999, module_dirs=None):
         self.host = host
         self.port = port
-        self.module_manager = ModuleManager()
+        self.module_dirs = module_dirs or ['modules']
+        self.module_manager = ModuleManager(self.module_dirs)
 
     def start(self):
-        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()
+        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)
 
     def handle_client(self, conn):
         with conn:
             while True:
-                data = conn.recv(1024)
-                if not data:
-                    break
-                command = json.loads(data.decode())
-                response = self.process_command(command)
-                conn.sendall(json.dumps(response).encode())
+                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())
 
     def upgrade(self):
-        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)
+        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)}"
 
     def process_command(self, command):
         action = command.get('action')
-        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}
+        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 __name__ == "__main__":
-    daemon = CoreDaemon()
-    daemon.start()
+    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_cli.py

@@ -1,6 +1,7 @@
 import cmd
 import socket
 import json
+import time
 import os
 
 class ModuleShell(cmd.Cmd):
@@ -13,17 +14,41 @@ class ModuleShell(cmd.Cmd):
         self.port = port
         self.update_commands()
 
-    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 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 update_commands(self):
-        response = self.send_command({'action': 'list_commands'})
-        if response['success']:
-            self.dynamic_commands = response['message']
-        else:
+        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)}")
             self.dynamic_commands = []
 
     def default(self, line):
@@ -31,10 +56,13 @@ class ModuleShell(cmd.Cmd):
         command = parts[0]
         args = ' '.join(parts[1:])
         response = self.send_command({'action': 'execute', 'command': command, 'args': args})
-        if response['success']:
-            print(response['message'])
+        if isinstance(response, dict):
+            if response.get('success'):
+                print(response['message'])
+            else:
+                print(f"Error: {response.get('message', 'Unknown error')}")
         else:
-            print(f"Error: {response['message']}")
+            print(f"Unexpected response: {response}")
 
     def completenames(self, text, *ignored):
         dotext = 'do_'+text
@@ -44,19 +72,32 @@ class ModuleShell(cmd.Cmd):
     def do_load(self, arg):
         """Load a module: load <module_name>"""
         response = self.send_command({'action': 'load', 'module': arg})
-        print(response['message'])
+        if isinstance(response, dict):
+            print(response.get('message', 'Unknown response'))
+        else:
+            print(f"Unexpected response: {response}")
         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})
-        print(response['message'])
+        if isinstance(response, dict):
+            print(response.get('message', 'Unknown response'))
+        else:
+            print(f"Unexpected response: {response}")
         self.update_commands()  # Update commands after unloading a module
 
     def do_list(self, arg):
         """List all loaded modules"""
         response = self.send_command({'action': 'list'})
-        modules = response['message']
+        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
+
         if modules:
             print("Loaded modules:")
             for module in modules:
@@ -77,17 +118,44 @@ 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'})
-        print(response['message'])
+        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
 
     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()