core_daemon/daemon_cli.py

import cmd
import socket
import json
import time
import os

class ModuleShell(cmd.Cmd):
    intro = "Welcome to the Module Shell. Type 'help' for commands."
    prompt = "(module_shell) "
    
    def __init__(self, host='localhost', port=9999):
        super().__init__()
        self.host = host
        self.port = port
        self.update_commands()

    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):
        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):
        parts = line.split()
        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'):
                print(response['message'])
            else:
                print(f"Error: {response.get('message', 'Unknown error')}")
        else:
            print(f"Unexpected response: {response}")

    def completenames(self, text, *ignored):
        dotext = 'do_'+text
        return [a[3:] for a in self.get_names() if a.startswith(dotext)] + \
               [cmd for cmd in self.dynamic_commands if cmd.startswith(text)]

    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}")
        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}")
        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:
                print(f"- {module}")
        else:
            print("No modules are currently loaded.")

    def do_list_available(self, arg):
        """List available modules in the modules folder"""
        modules_path = os.path.join(os.path.dirname(__file__), 'modules')
        available_modules = [f[:-3] for f in os.listdir(modules_path) if f.endswith('.py') and not f.startswith('__')]
        print("Available modules:")
        for module in available_modules:
            print(f"- {module}")

    def complete_load(self, text, line, begidx, endidx):
        modules_path = os.path.join(os.path.dirname(__file__), 'modules')
        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

    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()

if __name__ == "__main__":
    main()