core_daemon/core_daemon.py

165 lines
6.7 KiB
Python

import os
import time
import sys
import socket
import json
import threading
import logging
import traceback
import hashlib
from hook_manager import hook_manager
from module_manager import ModuleManager
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
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']
self.module_manager = ModuleManager(self.module_dirs)
self.script_path = sys.argv[0]
self.initial_hash = self.calculate_file_hash(self.script_path)
self.hook_manager = hook_manager
def calculate_file_hash(self, file_path):
sha256_hash = hashlib.sha256()
with open(file_path, "rb") as f:
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest()
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}")
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 check_for_updates(self):
logging.info("Checking for updates...")
script_modified = False
current_script_hash = self.calculate_file_hash(self.script_path)
if current_script_hash != self.initial_hash:
script_modified = True
modified_modules = self.module_manager.check_modules_modified()
if script_modified or modified_modules:
message = "Updates detected:\n"
if script_modified:
message += "- Core script has been modified\n"
if modified_modules:
message += f"- Modified modules: {', '.join(modified_modules)}\n"
return True, message
else:
return False, "No updates detected."
def upgrade(self):
update_available, message = self.check_for_updates()
if update_available:
logging.info(message)
# Here you would typically perform any necessary upgrade steps
# For this example, we'll just restart the script
os.execv(sys.executable, ['python'] + sys.argv)
return True, "Upgrade initiated. Restarting with new version."
else:
logging.info("No updates detected. Skipping upgrade.")
return False, "No upgrade needed. Current version is up to date."
def process_command(self, command):
action = command.get('action')
# Execute pre-command hooks
pre_command_results = self.hook_manager.execute_hook('pre_command', command)
for result in pre_command_results:
if result is False:
return {'success': False, 'message': "Command rejected by pre-command hook"}
try:
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':
success, modules = True, self.module_manager.list_modules()
message = modules
elif action == 'execute':
success, message = self.module_manager.execute_command(command.get('command'), command.get('args'))
elif action == 'upgrade':
success, message = self.upgrade()
elif action == 'list_commands':
success, commands = True, self.module_manager.list_commands()
message = commands
elif action == 'add_module_dir':
success, message = 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:
success, message = False, "Unknown command"
except Exception as e:
logging.error(f"Error processing command {action}: {str(e)}")
success, message = False, f"Error processing command: {str(e)}"
response = {'success': success, 'message': message}
# Execute post-command hooks
self.hook_manager.execute_hook('post_command', command, response)
return response
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.start()
except Exception as e:
logging.critical(f"Critical error in core daemon: {str(e)}")
logging.debug(traceback.format_exc())
sys.exit(1)