core_daemon/core_daemon.py

110 lines
4.1 KiB
Python

import os
import sys
import importlib
import socket
import json
import threading
# 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):
self.loaded_modules = {}
self.extra_commands = {}
def load_module(self, module_name):
try:
# 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)
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:
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."
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:
return True, self.extra_commands[command](args)
return False, "Command not found"
class CoreDaemon:
def __init__(self, host='localhost', port=9999):
self.host = host
self.port = port
self.module_manager = ModuleManager()
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()
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())
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)
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
else:
success, message = False, "Unknown command"
return {'success': success, 'message': message}
if __name__ == "__main__":
daemon = CoreDaemon()
daemon.start()