196 lines
8.3 KiB
Python
196 lines
8.3 KiB
Python
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)
|