311 lines
11 KiB
Python
311 lines
11 KiB
Python
|
"""
|
||
|
|
cmd2-like CLI framework using prompt_toolkit with:
|
||
|
|
|
||
|
|
- Color support
|
||
|
|
- Command history search
|
||
|
|
- Script execution
|
||
|
|
- Nested commands
|
||
|
|
- Session logging
|
||
|
|
"""
|
||
|
|
|
||
|
|
from prompt_toolkit import PromptSession, HTML
|
||
|
|
from prompt_toolkit.history import FileHistory
|
||
|
|
from prompt_toolkit.completion import Completer, Completion
|
||
|
|
from prompt_toolkit.lexers import PygmentsLexer
|
||
|
|
from prompt_toolkit.styles import Style
|
||
|
|
from pygments.lexers.shell import BashLexer
|
||
|
|
import inspect
|
||
|
|
import os
|
||
|
|
import datetime
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
|
||
|
|
# Custom color style
|
||
|
|
custom_style = Style.from_dict({
|
||
|
|
'prompt': 'ansicyan bold',
|
||
|
|
'command': 'ansigreen',
|
||
|
|
'error': 'ansired bold',
|
||
|
|
'warning': 'ansiyellow',
|
||
|
|
'info': 'ansiblue',
|
||
|
|
'success': 'ansigreen bold',
|
||
|
|
})
|
||
|
|
|
||
|
|
class CommandCompleter(Completer):
|
||
|
|
"""Provides autocompletion for commands and their arguments"""
|
||
|
|
def __init__(self, cli):
|
||
|
|
self.cli = cli
|
||
|
|
|
||
|
|
def get_completions(self, document, complete_event):
|
||
|
|
text = document.text_before_cursor.split()
|
||
|
|
|
||
|
|
# Command name completion
|
||
|
|
if len(text) == 0 or (len(text) == 1 and not document.text.endswith(' ')):
|
||
|
|
prefix = text[0] if text else ''
|
||
|
|
for cmd in self.cli.get_commands():
|
||
|
|
if cmd.startswith(prefix):
|
||
|
|
yield Completion(
|
||
|
|
cmd,
|
||
|
|
start_position=-len(prefix),
|
||
|
|
display=cmd,
|
||
|
|
display_meta=self.cli.get_command_help(cmd)
|
||
|
|
)
|
||
|
|
|
||
|
|
# Command argument completion
|
||
|
|
elif len(text) >= 1:
|
||
|
|
cmd = text[0]
|
||
|
|
arg_completer = getattr(self.cli, f'argcompleter_{cmd}', None)
|
||
|
|
if arg_completer:
|
||
|
|
args = ' '.join(text[1:])
|
||
|
|
for comp in arg_completer(args):
|
||
|
|
yield Completion(comp, start_position=-len(document.current_line))
|
||
|
|
|
||
|
|
class Cmd2LikeCLI:
|
||
|
|
"""Main CLI class with advanced features"""
|
||
|
|
def __init__(self):
|
||
|
|
self.session = PromptSession(
|
||
|
|
history=FileHistory('.cmd_history'),
|
||
|
|
completer=CommandCompleter(self),
|
||
|
|
lexer=PygmentsLexer(BashLexer),
|
||
|
|
complete_while_typing=True,
|
||
|
|
enable_history_search=True,
|
||
|
|
style=custom_style
|
||
|
|
)
|
||
|
|
self.prompt = HTML('<prompt>>>> </prompt>')
|
||
|
|
self.intro = self.colorize("Welcome to Advanced CLI (Type 'help' for commands)", "info")
|
||
|
|
self.logging_active = False
|
||
|
|
self.log_file = None
|
||
|
|
self.log_file_path = "cli_session.log"
|
||
|
|
|
||
|
|
def colorize(self, text, style_type):
|
||
|
|
"""Apply color to text based on style type"""
|
||
|
|
styles = {
|
||
|
|
"command": "<command>{}</command>",
|
||
|
|
"error": "<error>{}</error>",
|
||
|
|
"warning": "<warning>{}</warning>",
|
||
|
|
"info": "<info>{}</info>",
|
||
|
|
"success": "<success>{}</success>",
|
||
|
|
}
|
||
|
|
return HTML(styles.get(style_type, "{}").format(text))
|
||
|
|
|
||
|
|
def get_commands(self):
|
||
|
|
"""Retrieves all available commands (methods starting with 'do_')"""
|
||
|
|
return [
|
||
|
|
name[3:] for name, _ in inspect.getmembers(self, inspect.ismethod)
|
||
|
|
if name.startswith('do_')
|
||
|
|
]
|
||
|
|
|
||
|
|
def get_command_help(self, cmd):
|
||
|
|
"""Gets help text from command's docstring"""
|
||
|
|
method = getattr(self, f'do_{cmd}', None)
|
||
|
|
return method.__doc__.strip() if method and method.__doc__ else "No help available"
|
||
|
|
|
||
|
|
def start_logging(self):
|
||
|
|
"""Start session logging"""
|
||
|
|
try:
|
||
|
|
self.log_file = open(self.log_file_path, 'a')
|
||
|
|
self.logging_active = True
|
||
|
|
timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
|
self.log_file.write(f"\n\n--- Session started at {timestamp} ---\n")
|
||
|
|
print(self.colorize(f"Logging started to {self.log_file_path}", "success"))
|
||
|
|
except Exception as e:
|
||
|
|
print(self.colorize(f"Error starting logging: {e}", "error"))
|
||
|
|
|
||
|
|
def stop_logging(self):
|
||
|
|
"""Stop session logging"""
|
||
|
|
if self.logging_active:
|
||
|
|
try:
|
||
|
|
timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
|
self.log_file.write(f"\n--- Session ended at {timestamp} ---\n")
|
||
|
|
self.log_file.close()
|
||
|
|
self.logging_active = False
|
||
|
|
print(self.colorize("Logging stopped", "success"))
|
||
|
|
except Exception as e:
|
||
|
|
print(self.colorize(f"Error stopping logging: {e}", "error"))
|
||
|
|
|
||
|
|
def log_command(self, command):
|
||
|
|
"""Log command to file if logging is active"""
|
||
|
|
if self.logging_active:
|
||
|
|
try:
|
||
|
|
timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
|
self.log_file.write(f"[{timestamp}] {command}\n")
|
||
|
|
except Exception as e:
|
||
|
|
print(self.colorize(f"Logging error: {e}", "error"))
|
||
|
|
|
||
|
|
def cmdloop(self):
|
||
|
|
"""Main REPL loop with nested command support"""
|
||
|
|
print(self.intro)
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
user_input = self.session.prompt(self.prompt)
|
||
|
|
# Process nested commands (separated by semicolons)
|
||
|
|
for cmd in user_input.split(';'):
|
||
|
|
cmd = cmd.strip()
|
||
|
|
if cmd:
|
||
|
|
self.onecmd(cmd)
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
print(self.colorize("\nUse 'exit' to quit", "warning"))
|
||
|
|
except EOFError:
|
||
|
|
self.do_exit("")
|
||
|
|
break
|
||
|
|
|
||
|
|
def onecmd(self, line):
|
||
|
|
"""Executes a single command"""
|
||
|
|
line = line.strip()
|
||
|
|
if not line:
|
||
|
|
return
|
||
|
|
|
||
|
|
# Log command before execution
|
||
|
|
self.log_command(line)
|
||
|
|
|
||
|
|
parts = line.split(maxsplit=1)
|
||
|
|
cmd_name = parts[0]
|
||
|
|
args = parts[1] if len(parts) > 1 else ""
|
||
|
|
|
||
|
|
if cmd_name == "help":
|
||
|
|
self.do_help(args)
|
||
|
|
return
|
||
|
|
|
||
|
|
method_name = f"do_{cmd_name}"
|
||
|
|
method = getattr(self, method_name, None)
|
||
|
|
|
||
|
|
if method:
|
||
|
|
try:
|
||
|
|
method(args)
|
||
|
|
except Exception as e:
|
||
|
|
print(self.colorize(f"Error executing command: {e}", "error"))
|
||
|
|
else:
|
||
|
|
print(self.colorize(f"Unknown command: {cmd_name}", "error"))
|
||
|
|
|
||
|
|
def run_script(self, file_path):
|
||
|
|
"""Execute commands from a script file"""
|
||
|
|
try:
|
||
|
|
with open(file_path, 'r') as f:
|
||
|
|
for line in f:
|
||
|
|
line = line.strip()
|
||
|
|
if line and not line.startswith('#'): # Skip empty lines and comments
|
||
|
|
print(self.colorize(f"Executing: {line}", "info"))
|
||
|
|
self.onecmd(line)
|
||
|
|
print(self.colorize(f"Script '{file_path}' executed successfully", "success"))
|
||
|
|
except FileNotFoundError:
|
||
|
|
print(self.colorize(f"Script file not found: {file_path}", "error"))
|
||
|
|
except Exception as e:
|
||
|
|
print(self.colorize(f"Error executing script: {e}", "error"))
|
||
|
|
|
||
|
|
# ------------ Built-in commands ------------
|
||
|
|
def do_help(self, arg):
|
||
|
|
"""Show help for commands"""
|
||
|
|
if arg:
|
||
|
|
# Command-specific help
|
||
|
|
help_text = self.get_command_help(arg)
|
||
|
|
print(f"{self.colorize(arg, 'command')}: {help_text}")
|
||
|
|
else:
|
||
|
|
# List all commands
|
||
|
|
print(self.colorize("\nAvailable commands:", "info"))
|
||
|
|
for cmd in self.get_commands():
|
||
|
|
print(f" {self.colorize(cmd, 'command'):15} {self.get_command_help(cmd)}")
|
||
|
|
print()
|
||
|
|
|
||
|
|
def do_echo(self, arg):
|
||
|
|
"""Echo back the input"""
|
||
|
|
print(f"Echo: {arg}")
|
||
|
|
|
||
|
|
def do_exit(self, arg):
|
||
|
|
"""Exit the application"""
|
||
|
|
self.stop_logging()
|
||
|
|
print(self.colorize("Goodbye!", "success"))
|
||
|
|
raise SystemExit
|
||
|
|
|
||
|
|
def do_ls(self, arg):
|
||
|
|
"""List directory contents"""
|
||
|
|
path = arg.strip() if arg else "."
|
||
|
|
try:
|
||
|
|
files = os.listdir(path)
|
||
|
|
for f in files:
|
||
|
|
full_path = os.path.join(path, f)
|
||
|
|
if os.path.isdir(full_path):
|
||
|
|
print(self.colorize(f + "/", "info"))
|
||
|
|
else:
|
||
|
|
print(f)
|
||
|
|
except FileNotFoundError:
|
||
|
|
print(self.colorize(f"Directory not found: {path}", "error"))
|
||
|
|
|
||
|
|
def do_run_script(self, arg):
|
||
|
|
"""Execute commands from a script file"""
|
||
|
|
if not arg:
|
||
|
|
print(self.colorize("Please specify a script file", "warning"))
|
||
|
|
return
|
||
|
|
self.run_script(arg.strip())
|
||
|
|
|
||
|
|
def do_log_start(self, arg):
|
||
|
|
"""Start session logging"""
|
||
|
|
if arg:
|
||
|
|
self.log_file_path = arg.strip()
|
||
|
|
self.start_logging()
|
||
|
|
|
||
|
|
def do_log_stop(self, arg):
|
||
|
|
"""Stop session logging"""
|
||
|
|
self.stop_logging()
|
||
|
|
|
||
|
|
def do_save_history(self, arg):
|
||
|
|
"""Save command history to file"""
|
||
|
|
file_path = arg.strip() if arg else "command_history.txt"
|
||
|
|
try:
|
||
|
|
with open('.cmd_history', 'r') as src, open(file_path, 'w') as dest:
|
||
|
|
dest.write(f"Command History - Saved at {datetime.datetime.now()}\n\n")
|
||
|
|
dest.writelines(src.readlines())
|
||
|
|
print(self.colorize(f"History saved to {file_path}", "success"))
|
||
|
|
except Exception as e:
|
||
|
|
print(self.colorize(f"Error saving history: {e}", "error"))
|
||
|
|
|
||
|
|
def do_system(self, arg):
|
||
|
|
"""Execute a system shell command"""
|
||
|
|
if not arg:
|
||
|
|
print(self.colorize("Please specify a command to execute", "warning"))
|
||
|
|
return
|
||
|
|
|
||
|
|
try:
|
||
|
|
result = subprocess.run(
|
||
|
|
arg,
|
||
|
|
shell=True,
|
||
|
|
capture_output=True,
|
||
|
|
text=True
|
||
|
|
)
|
||
|
|
if result.stdout:
|
||
|
|
print(result.stdout)
|
||
|
|
if result.stderr:
|
||
|
|
print(self.colorize(result.stderr, "error"))
|
||
|
|
print(self.colorize(f"Command exited with code: {result.returncode}", "info"))
|
||
|
|
except Exception as e:
|
||
|
|
print(self.colorize(f"Error executing system command: {e}", "error"))
|
||
|
|
|
||
|
|
# ------------ Custom argument completers ------------
|
||
|
|
def argcompleter_ls(self, arg):
|
||
|
|
"""Directory completer for ls command"""
|
||
|
|
current_dir = os.path.dirname(arg) if arg else "."
|
||
|
|
base_name = os.path.basename(arg) if arg else ""
|
||
|
|
|
||
|
|
try:
|
||
|
|
return [
|
||
|
|
f.name for f in os.scandir(current_dir)
|
||
|
|
if f.name.startswith(base_name)
|
||
|
|
]
|
||
|
|
except FileNotFoundError:
|
||
|
|
return []
|
||
|
|
|
||
|
|
def argcompleter_run_script(self, arg):
|
||
|
|
"""File completer for run_script command"""
|
||
|
|
current_dir = os.path.dirname(arg) if arg else "."
|
||
|
|
base_name = os.path.basename(arg) if arg else ""
|
||
|
|
|
||
|
|
try:
|
||
|
|
return [
|
||
|
|
f.name for f in os.scandir(current_dir)
|
||
|
|
if f.is_file() and f.name.startswith(base_name)
|
||
|
|
]
|
||
|
|
except FileNotFoundError:
|
||
|
|
return []
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
cli = Cmd2LikeCLI()
|
||
|
|
cli.cmdloop()
|