core_daemon/core_daemon.py

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)