core_daemon/daemon_cli.py
2024-08-06 22:38:33 +03:00

164 lines
6.5 KiB
Python

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