Now we can shutdown the daemon from the cli.
This commit is contained in:
kalzu rekku
parent
3b919958af
commit
2317fa4352
176
core_daemon.py
176
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)
|
||||
|
||||
100
daemon_cli.py
100
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()
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user