From 2317fa43520029ec39824075a2ec4c7787bd1525 Mon Sep 17 00:00:00 2001 From: kalzu rekku Date: Tue, 6 Aug 2024 22:38:33 +0300 Subject: [PATCH] Now we can shutdown the daemon from the cli. --- core_daemon.py | 176 ++++++++++++++++++++++++++++++++----------------- daemon_cli.py | 100 +++++++++++++++++++++++----- 2 files changed, 201 insertions(+), 75 deletions(-) diff --git a/core_daemon.py b/core_daemon.py index f18846c..5157106 100644 --- a/core_daemon.py +++ b/core_daemon.py @@ -1,13 +1,15 @@ 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, module_dirs): @@ -21,40 +23,55 @@ class ModuleManager: 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): - if new_dir not in self.module_dirs: - self.module_dirs.append(new_dir) - self._update_sys_path() - return True, f"Added module directory: {new_dir}" - return False, f"Module directory already exists: {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'{dir}.{module_name}') + 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): @@ -65,63 +82,99 @@ 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, module_dirs=None): self.host = host self.port = port - self.module_dirs = module_dirs or ['modules'] # Default to 'modules' if not specified + 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 - elif action == 'add_module_dir': - success, message = self.module_manager.add_module_dir(command.get('dir')) - 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__": import argparse @@ -133,5 +186,10 @@ if __name__ == "__main__": args = parser.parse_args() - daemon = CoreDaemon(host=args.host, port=args.port, module_dirs=args.module_dirs) - daemon.start() + 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) diff --git a/daemon_cli.py b/daemon_cli.py index 1da63e1..46b61d3 100644 --- a/daemon_cli.py +++ b/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 """ 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 """ 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 """ + 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()