Compare commits

..

No commits in common. "6a86a5bb4ab56ce75d5591b431f4c60416bae856" and "2a64d79c5185eebdf9b64beaa1331db2674b6640" have entirely different histories.

3 changed files with 85 additions and 237 deletions

View File

@ -63,6 +63,8 @@ See `modules/extra_commands.py` and `modules/http_service.py` for an example of
# Todo
- Create module to add auth for the cli connections...
- Create module for module to module communication...
- Maybe add ability to search for modules in directories set by a command
- Create module to add auth for the cli connections...
- Add error handling to core_daemon and daemon_cli...

View File

@ -1,64 +1,35 @@
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')
# Add the modules folder to the Python path
modules_path = os.path.join(os.path.dirname(__file__), 'modules')
sys.path.append(modules_path)
class ModuleManager:
def __init__(self, module_dirs):
self.module_dirs = module_dirs
def __init__(self):
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}')
# Try to import from the modules folder
module = importlib.import_module(f'modules.{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."
return True, f"Module '{module_name}' loaded and initialized successfully."
except ImportError as e:
return False, f"Error: Unable to load module '{module_name}'. {str(e)}"
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()
@ -67,11 +38,7 @@ class ModuleManager:
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):
@ -82,114 +49,61 @@ class ModuleManager:
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 True, self.extra_commands[command](args)
return False, "Command not found"
class CoreDaemon:
def __init__(self, host='localhost', port=9999, module_dirs=None):
def __init__(self, host='localhost', port=9999):
self.host = host
self.port = port
self.module_dirs = module_dirs or ['modules']
self.module_manager = ModuleManager(self.module_dirs)
self.module_manager = ModuleManager()
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}")
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()
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:
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)
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'))
success, message = self.module_manager.load_module(command.get('module'))
elif action == 'unload':
return self.module_manager.unload_module(command.get('module'))
success, message = self.module_manager.unload_module(command.get('module'))
elif action == 'list':
modules = self.module_manager.list_modules()
return True, modules
success, message = True, modules
elif action == 'execute':
return self.module_manager.execute_command(command.get('command'), command.get('args'))
success, message = self.module_manager.execute_command(command.get('command'), command.get('args'))
elif action == 'upgrade':
return self.upgrade()
self.upgrade()
success, message = True, "Upgrade initiated"
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
success, message = True, commands
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)}"
success, message = False, "Unknown command"
return {'success': success, 'message': message}
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 = CoreDaemon()
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,7 +1,6 @@
import cmd
import socket
import json
import time
import os
class ModuleShell(cmd.Cmd):
@ -14,41 +13,17 @@ class ModuleShell(cmd.Cmd):
self.port = port
self.update_commands()
def send_command(self, command, retries=3, delay=2):
for attempt in range(retries):
try:
def send_command(self, command):
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."}
return json.loads(s.recv(1024).decode())
def update_commands(self):
try:
response = self.send_command({'action': 'list_commands'})
if isinstance(response, dict) and response.get('success'):
if response['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):
@ -56,13 +31,10 @@ class ModuleShell(cmd.Cmd):
command = parts[0]
args = ' '.join(parts[1:])
response = self.send_command({'action': 'execute', 'command': command, 'args': args})
if isinstance(response, dict):
if response.get('success'):
if response['success']:
print(response['message'])
else:
print(f"Error: {response.get('message', 'Unknown error')}")
else:
print(f"Unexpected response: {response}")
print(f"Error: {response['message']}")
def completenames(self, text, *ignored):
dotext = 'do_'+text
@ -72,32 +44,19 @@ class ModuleShell(cmd.Cmd):
def do_load(self, arg):
"""Load a module: load <module_name>"""
response = self.send_command({'action': 'load', 'module': arg})
if isinstance(response, dict):
print(response.get('message', 'Unknown response'))
else:
print(f"Unexpected response: {response}")
print(response['message'])
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})
if isinstance(response, dict):
print(response.get('message', 'Unknown response'))
else:
print(f"Unexpected response: {response}")
print(response['message'])
self.update_commands() # Update commands after unloading a module
def do_list(self, arg):
"""List all loaded modules"""
response = self.send_command({'action': 'list'})
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:
@ -118,44 +77,17 @@ 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'})
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
print(response['message'])
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()