Now we can shutdown the daemon from the cli.

This commit is contained in:
kalzu rekku 2024-08-06 22:38:33 +03:00
parent 3b919958af
commit 2317fa4352
2 changed files with 201 additions and 75 deletions

View File

@ -1,13 +1,15 @@
import os import os
import time
import sys import sys
import importlib import importlib
import socket import socket
import json import json
import threading import threading
import logging
import traceback
# Add the modules folder to the Python path # Configure logging
modules_path = os.path.join(os.path.dirname(__file__), 'modules') logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
sys.path.append(modules_path)
class ModuleManager: class ModuleManager:
def __init__(self, module_dirs): def __init__(self, module_dirs):
@ -21,31 +23,42 @@ class ModuleManager:
full_path = os.path.abspath(dir) full_path = os.path.abspath(dir)
if full_path not in sys.path: if full_path not in sys.path:
sys.path.append(full_path) sys.path.append(full_path)
logging.info(f"Added {full_path} to sys.path")
def add_module_dir(self, new_dir): def add_module_dir(self, new_dir):
if new_dir not in self.module_dirs: try:
self.module_dirs.append(new_dir) 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() self._update_sys_path()
return True, f"Added module directory: {new_dir}" return True, f"Added module directory: {full_path}"
return False, f"Module directory already exists: {new_dir}" 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): def load_module(self, module_name):
for dir in self.module_dirs: for dir in self.module_dirs:
try: 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 self.loaded_modules[module_name] = module
if hasattr(module, 'initialize'): if hasattr(module, 'initialize'):
module.initialize() module.initialize()
if hasattr(module, 'get_commands'): if hasattr(module, 'get_commands'):
new_commands = module.get_commands() new_commands = module.get_commands()
self.extra_commands.update(new_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}." return True, f"Module '{module_name}' loaded and initialized successfully from {dir}."
except ImportError: except ImportError:
continue 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." return False, f"Error: Unable to load module '{module_name}' from any of the module directories."
def unload_module(self, module_name): def unload_module(self, module_name):
if module_name in self.loaded_modules: if module_name in self.loaded_modules:
try:
module = self.loaded_modules[module_name] module = self.loaded_modules[module_name]
if hasattr(module, 'shutdown'): if hasattr(module, 'shutdown'):
module.shutdown() module.shutdown()
@ -54,7 +67,11 @@ class ModuleManager:
for cmd in commands_to_remove: for cmd in commands_to_remove:
self.extra_commands.pop(cmd, None) self.extra_commands.pop(cmd, None)
del self.loaded_modules[module_name] del self.loaded_modules[module_name]
logging.info(f"Module '{module_name}' unloaded successfully.")
return True, f"Module '{module_name}' unloaded and shut down." 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." return False, f"Module '{module_name}' is not loaded."
def list_modules(self): def list_modules(self):
@ -65,63 +82,99 @@ class ModuleManager:
def execute_command(self, command, args): def execute_command(self, command, args):
if command in self.extra_commands: 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" return False, "Command not found"
class CoreDaemon: class CoreDaemon:
def __init__(self, host='localhost', port=9999, module_dirs=None): def __init__(self, host='localhost', port=9999, module_dirs=None):
self.host = host self.host = host
self.port = port 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) self.module_manager = ModuleManager(self.module_dirs)
def start(self): def start(self):
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: 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.bind((self.host, self.port))
s.listen() s.listen()
print(f"Core daemon listening on {self.host}:{self.port}") logging.info(f"Core daemon listening on {self.host}:{self.port}")
while True: while True:
conn, addr = s.accept() conn, addr = s.accept()
threading.Thread(target=self.handle_client, args=(conn,)).start() 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): def handle_client(self, conn):
with conn: with conn:
while True: while True:
try:
data = conn.recv(1024) data = conn.recv(1024)
if not data: if not data:
break break
command = json.loads(data.decode()) command = json.loads(data.decode())
response = self.process_command(command) response = self.process_command(command)
conn.sendall(json.dumps(response).encode()) 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): def upgrade(self):
print("Upgrading core daemon...") logging.info("Upgrading core daemon...")
try:
# Here you would typically download or copy the new version # Here you would typically download or copy the new version
# For this example, we'll just restart the script # For this example, we'll just restart the script
os.execv(sys.executable, ['python'] + sys.argv) 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): def process_command(self, command):
action = command.get('action') action = command.get('action')
try:
if action == 'load': if action == 'load':
success, message = self.module_manager.load_module(command.get('module')) return self.module_manager.load_module(command.get('module'))
elif action == 'unload': elif action == 'unload':
success, message = self.module_manager.unload_module(command.get('module')) return self.module_manager.unload_module(command.get('module'))
elif action == 'list': elif action == 'list':
modules = self.module_manager.list_modules() modules = self.module_manager.list_modules()
success, message = True, modules return True, modules
elif action == 'execute': elif action == 'execute':
success, message = self.module_manager.execute_command(command.get('command'), command.get('args')) return self.module_manager.execute_command(command.get('command'), command.get('args'))
elif action == 'upgrade': elif action == 'upgrade':
self.upgrade() return self.upgrade()
success, message = True, "Upgrade initiated"
elif action == 'list_commands': elif action == 'list_commands':
commands = self.module_manager.list_commands() commands = self.module_manager.list_commands()
success, message = True, commands return True, commands
elif action == 'add_module_dir': elif action == 'add_module_dir':
success, message = self.module_manager.add_module_dir(command.get('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: else:
success, message = False, "Unknown command" return False, "Unknown command"
return {'success': success, 'message': message} 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__": if __name__ == "__main__":
import argparse import argparse
@ -133,5 +186,10 @@ if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
try:
daemon = CoreDaemon(host=args.host, port=args.port, module_dirs=args.module_dirs) daemon = CoreDaemon(host=args.host, port=args.port, module_dirs=args.module_dirs)
daemon.start() daemon.start()
except Exception as e:
logging.critical(f"Critical error in core daemon: {str(e)}")
logging.debug(traceback.format_exc())
sys.exit(1)

View File

@ -1,6 +1,7 @@
import cmd import cmd
import socket import socket
import json import json
import time
import os import os
class ModuleShell(cmd.Cmd): class ModuleShell(cmd.Cmd):
@ -13,17 +14,41 @@ class ModuleShell(cmd.Cmd):
self.port = port self.port = port
self.update_commands() self.update_commands()
def send_command(self, command): 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: 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.connect((self.host, self.port))
s.sendall(json.dumps(command).encode()) s.sendall(json.dumps(command).encode())
return json.loads(s.recv(1024).decode()) 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): def update_commands(self):
try:
response = self.send_command({'action': 'list_commands'}) response = self.send_command({'action': 'list_commands'})
if response['success']: if isinstance(response, dict) and response.get('success'):
self.dynamic_commands = response['message'] self.dynamic_commands = response['message']
elif isinstance(response, list):
self.dynamic_commands = response
else: 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 = [] self.dynamic_commands = []
def default(self, line): def default(self, line):
@ -31,10 +56,13 @@ class ModuleShell(cmd.Cmd):
command = parts[0] command = parts[0]
args = ' '.join(parts[1:]) args = ' '.join(parts[1:])
response = self.send_command({'action': 'execute', 'command': command, 'args': args}) response = self.send_command({'action': 'execute', 'command': command, 'args': args})
if response['success']: if isinstance(response, dict):
if response.get('success'):
print(response['message']) print(response['message'])
else: else:
print(f"Error: {response['message']}") print(f"Error: {response.get('message', 'Unknown error')}")
else:
print(f"Unexpected response: {response}")
def completenames(self, text, *ignored): def completenames(self, text, *ignored):
dotext = 'do_'+text dotext = 'do_'+text
@ -44,19 +72,32 @@ class ModuleShell(cmd.Cmd):
def do_load(self, arg): def do_load(self, arg):
"""Load a module: load <module_name>""" """Load a module: load <module_name>"""
response = self.send_command({'action': 'load', 'module': arg}) 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 self.update_commands() # Update commands after loading a new module
def do_unload(self, arg): def do_unload(self, arg):
"""Unload a module: unload <module_name>""" """Unload a module: unload <module_name>"""
response = self.send_command({'action': 'unload', 'module': arg}) 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 self.update_commands() # Update commands after unloading a module
def do_list(self, arg): def do_list(self, arg):
"""List all loaded modules""" """List all loaded modules"""
response = self.send_command({'action': 'list'}) response = self.send_command({'action': 'list'})
if isinstance(response, dict) and response.get('success'):
modules = response['message'] modules = response['message']
elif isinstance(response, list):
modules = response
else:
print("Unexpected response format when listing modules.")
return
if modules: if modules:
print("Loaded modules:") print("Loaded modules:")
for module in 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('__')] 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)] 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): def do_upgrade(self, arg):
"""Upgrade the core daemon""" """Upgrade the core daemon"""
response = self.send_command({'action': 'upgrade'}) 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): def do_exit(self, arg):
"""Exit the shell""" """Exit the shell"""
print("Exiting...") print("Exiting...")
return True 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(): def main():
ModuleShell().cmdloop() ModuleShell().cmdloop()