core_daemon/core_daemon.py

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

class ModuleManager:
    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):
        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:
            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):
        return list(self.loaded_modules.keys())

    def list_commands(self):
        return list(self.extra_commands.keys())

    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 False, "Command not found"

class CoreDaemon:
    def __init__(self, host='localhost', port=9999, module_dirs=None):
        self.host = host
        self.port = port
        self.module_dirs = module_dirs or ['modules']
        self.module_manager = ModuleManager(self.module_dirs)

    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)

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

    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)}"

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