class CompCommand(Command): def __init__(self, name, topic, brief, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief) super(CompCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: self.parser.parse_args(args) except CommandShortCircuit as e: return e.code component_info = shell._component_info # read print_header from args ns = MainSectionWizard(component_info).run(shell, print_header=True) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) component_info.load_from_json(out_dict) else: return EXIT_CODE return 0
class ExitCommand(Command): ''' Exit the active shell. ''' def __init__(self, name='exit', topic='shell', brief='exit the shell', **kwargs): self.parser = PypsiArgParser( prog=name, description=brief ) self.parser.add_argument( 'code', metavar='CODE', action='store', default=0, type=int, nargs='?', help='exit code to return to parent process' ) super(ExitCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code raise SystemExit(ns.code)
class ExitCommand(Command): ''' Exit the active shell. ''' def __init__(self, name='exit', topic='shell', brief='exit the shell', **kwargs): self.parser = PypsiArgParser(prog=name, description=brief) self.parser.add_argument('code', metavar='CODE', action='store', default=0, type=int, nargs='?', help='exit code to return to parent process') super(ExitCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code raise SystemExit(ns.code)
class ShellCommand(Command): def __init__(self, topic='proxy', name='shell', brief=HelpTopic, **kwargs): self.parser = PypsiArgParser( prog=name, description=brief, usage=ShellCmdUsage ) super(ShellCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def run(self, shell, args): if not self.check_connected(shell): self.error("shell", "Not currently connected to minion") return 0 print(f'Starting interactive shell with: {shell.connect}') minionshell = MinionShell(shell.connect, shell) minionshell.cmdloop() print(f'Closed interactive shell with: {shell.connect}') def check_connected(self, shell): minion = shell.connect if minion is not None: return True
class SaveLoadCommand(Command): def __init__(self, name, topic, brief, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief) super(SaveLoadCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if self.name == CommandName.SAVE: w = SaveLoadWizard("Save Component", shell._filename) ns = w.run(shell, True) if ns: if ns.path is not None: component_info = shell._component_info with io.open(ns.path, mode='w', encoding="utf-8") as f: json.dump(component_info, f, indent=4, cls=ComponentInfoEncoder) print("Component was successfully saved: {}".format( ns.path)) shell._filename = ns.path else: print("Abort saving component") elif self.name == CommandName.LOAD: w = SaveLoadWizard("Load Component", None) ns = w.run(shell, True) if ns: with io.open(ns.path, encoding="utf-8") as f: try: json_dict = json.load(f) shell._component_info.load_from_json(json_dict) except Exception as e: print( "Can not load file {}. Not a json file?. Error: {}" .format(ns.path, str(e))) return 1 shell._filename = ns.path print("Component was successfully loaded: {}".format(ns.path)) else: pass return 0
class XArgsCommand(Command): ''' Execute a command for each line of input from :data:`sys.stdin`. ''' def __init__(self, name='xargs', topic='shell', brief='build and execute command lines from stdin', **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=XArgsUsage.format(name=name)) self.parser.add_argument('-I', default='{}', action='store', metavar='REPSTR', help='string token to replace', dest='token') self.parser.add_argument('command', nargs=argparse.REMAINDER, help="command to execute", metavar='COMMAND') super(XArgsCommand, self).__init__(name=name, topic=topic, usage=self.parser.format_help(), brief=brief, **kwargs) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if not ns.command: self.error(shell, "missing command") return 1 base = ' '.join( ['"{}"'.format(c.replace('"', '\\"')) for c in ns.command]) for line in sys.stdin: cmd = base.replace(ns.token, line.strip()) shell.execute(cmd) return 0
class XArgsCommand(Command): ''' Execute a command for each line of input from :data:`sys.stdin`. ''' def __init__(self, name='xargs', topic='shell', brief='build and execute command lines from stdin', **kwargs): self.parser = PypsiArgParser( prog=name, description=brief, usage=XArgsUsage.format(name=name) ) self.parser.add_argument( '-I', default='{}', action='store', metavar='REPSTR', help='string token to replace', dest='token' ) self.parser.add_argument( 'command', nargs=argparse.REMAINDER, help="command to execute", metavar='COMMAND' ) super(XArgsCommand, self).__init__( name=name, topic=topic, usage=self.parser.format_help(), brief=brief, **kwargs ) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if not ns.command: self.error(shell, "missing command") return 1 base = ' '.join([ '"{}"'.format(c.replace('"', '\\"')) for c in ns.command ]) for line in sys.stdin: cmd = base.replace(ns.token, line.strip()) shell.execute(cmd) return 0
class ShowCommand(Command): def __init__(self, name, topic, brief, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief) super(ShowCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: self.parser.parse_args(args) except CommandShortCircuit as e: return e.code component_info = shell._component_info print(json.dumps(component_info, indent=4, cls=ComponentInfoEncoder)) return 0
class DisconnectCommand(Command): def __init__(self, topic='proxy', name='disconnect', brief=HelpTopic, **kwargs): self.parser = PypsiArgParser( prog=name, description=brief, usage=DisconnectCmdUsage ) self.parser.add_argument( 'name', help='name of minion', metavar="NAME", completer=self.complete_minions ) super(DisconnectCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if not ns.name: pass elif ns.name == shell.connect: shell.connect = None shell.remove_prompt_connect() print(f"Successfully disconnected from minion: {ns.name}") else: self.error(shell,f"master is not currently connect to {ns.name}") def complete(self, shell, args, prefix): if len(args) == 1: return self.complete_minions(shell, args, prefix) def complete_minions(self, shell, args, prefix): return sorted( [name for name in shell.minions if name.startswith(prefix)] )
class EchoCommand(Command): ''' Prints text to the screen. ''' def __init__(self, name='echo', topic='shell', brief='print a line of text', **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=EchoCmdUsage) self.parser.add_argument('message', help='message to print', nargs=argparse.REMAINDER, metavar="MESSAGE") self.parser.add_argument('-n', '--nolf', help="don't print newline character", action='store_true') super(EchoCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code tail = '' if ns.nolf else '\n' print(' '.join(ns.message), sep='', end=tail) return 0
class ListCommand(Command): def __init__(self, topic='proxy', brief=HelpTopic, name='list', **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=ListCmdUsage) super(ListCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): print("\nConfigured minions:\n") table = Table(columns=(Column('Name', Column.Shrink), Column('IP Address', Column.Shrink), Column('Plugins', Column.Grow)), spacing=4) for name in shell.minions: table.append(name, shell.minions[name]['ip'], obj_str(shell.minions[name]['plugins'])) table.write(sys.stdout) print()
class EchoCommand(Command): ''' Prints text to the screen. ''' def __init__(self, name='echo', topic='shell', brief='print a line of text', **kwargs): self.parser = PypsiArgParser( prog=name, description=brief, usage=EchoCmdUsage ) self.parser.add_argument( 'message', help='message to print', nargs=argparse.REMAINDER, metavar="MESSAGE" ) self.parser.add_argument( '-n', '--nolf', help="don't print newline character", action='store_true' ) super(EchoCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code tail = '' if ns.nolf else '\n' print(' '.join(ns.message), sep='', end=tail) return 0
class TailCommand(Command): ''' Displays the last N lines of a file to the screen. ''' def __init__(self, name='tail', topic='shell', brief='display the last lines of a file', **kwargs): self.parser = PypsiArgParser( prog=name, description=brief, usage=TailCmdUsage ) self.parser.add_argument( 'input_file', help='file to display', metavar="FILE" ) self.parser.add_argument( '-n', '--lines', metavar="N", type=int, default=10, help="number of lines to display" ) self.parser.add_argument( '-f', '--follow', help="continue to output as file grows", action='store_true' ) super(TailCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code # check for valid input file if not os.path.isfile(ns.input_file): self.error(shell, "invalid file path: ", ns.input_file, "\n") return -1 # print the last N lines last_lines = self.tail(ns.input_file, ns.lines) for line in last_lines: print(line) print() # continue to follow the file and display new content if ns.follow: self.follow_file(ns.input_file) return 0 def complete(self, shell, args, prefix): return path_completer(args[-1]) def tail(self, fname, lines=10, block_size=1024): data = [] blocks = -1 num_lines = 0 with open(fname) as fp: # seek to the end and get the number of bytes in the file fp.seek(0, 2) num_bytes = fp.tell() bytes_left = num_bytes while num_lines < lines and bytes_left > 0: if bytes_left - block_size > 0: # Seek back a block_size fp.seek(num_bytes - (blocks * block_size)) # read data from file data.insert(0, fp.read(block_size)) else: # jump back to the beginning fp.seek(0, 0) # read data data.insert(0, fp.read(num_bytes)) num_lines = data[0].count('\n') bytes_left -= block_size blocks -= 1 return ''.join(data).splitlines()[-lines:] def follow_file(self, fname): with open(fname) as fp: # jump to the end of the file fp.seek(0, 2) try: while 1: where = fp.tell() line = fp.readline() if not line: time.sleep(1) fp.seek(where) else: print(line, end='', flush=True) except (KeyboardInterrupt): print() return 0 return 0
class TipCommand(Command): def __init__(self, name='tip', brief='print shell tips', topic='shell', tips=None, motd=None, vars=None, **kwargs): self.tips = tips or [] self.motd = motd or '' self.vars = vars or {} self.rand = random.Random() self.rand.seed() self.parser = PypsiArgParser(prog=name, description=brief) self.parser.add_argument('-m', '--motd', action='store_true', help="print the message of the day") super(TipCommand, self).__init__(name=name, brief=brief, topic=topic, usage=self.parser.format_help(), **kwargs) def load_tips(self, path): fp = safe_open(path, 'r') tip = [] for line in fp.readlines(): line = line.rstrip() if line: tip.append(line) elif tip: self.tips.append(' '.join(tip)) tip = [] if tip: self.tips.append(' '.join(tip)) fp.close() def load_motd(self, path): fp = safe_open(path, 'r') self.motd = fp.read().rstrip() fp.close() def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = None if ns.motd: rc = self.print_motd(shell) else: rc = self.print_random_tip(shell) return rc def print_random_tip(self, shell, header=True): if not self.tips: self.error(shell, "no tips available") return -1 i = self.rand.randrange(len(self.tips)) if header: title = "Tip #{}\n".format(i + 1) title += '-' * len(title) print(AnsiCodes.green, title, AnsiCodes.reset, sep='') try: cnt = sys.stdout.ansi_format(self.tips[i], **self.vars) except: cnt = self.tips[i] print(cnt) return 0 def print_motd(self, shell): if not self.motd: self.error(shell, "no motd available") return -1 try: cnt = sys.stdout.ansi_format(self.motd, **self.vars) except: cnt = self.motd print(AnsiCodes.green, "Message of the Day".center(shell.width), '\n', AnsiCodes.green, '>' * shell.width, "\n", AnsiCodes.reset, cnt, '\n', AnsiCodes.green, "<" * shell.width, "\n", AnsiCodes.reset, sep='') return 0
class MacroCommand(BlockCommand): ''' Record and execute command macros. Macros can consist of one or more command statements. Macros are comparable to Bash functions. Once a macro has been recorded, a new command is registered in the shell as an instance of a :class:`Macro`. This command requires the :class:`pypsi.plugins.block.BlockPlugin` plugin. ''' def __init__(self, name='macro', topic='shell', brief="manage registered macros", macros=None, **kwargs): self.parser = PypsiArgParser( prog=name, description=brief, usage=MacroCmdUsage ) self.parser.add_argument( '-l', '--list', help='list all macros', action='store_true' ) self.parser.add_argument( '-d', '--delete', help='delete macro', metavar='NAME', completer=self.complete_macros ) self.parser.add_argument( '-s', '--show', help='print macro body', metavar='NAME', completer=self.complete_macros ) self.parser.add_argument( 'name', help='macro name', nargs='?', metavar='NAME' ) super(MacroCommand, self).__init__( name=name, usage=self.parser.format_help(), brief=brief, topic=topic, **kwargs ) self.base_macros = macros or {} self.macro_name = None def complete_macros(self, shell, args, prefix): # pylint: disable=unused-argument # returns a list of macro names in the current shell return list(shell.ctx.macros.keys()) def complete(self, shell, args, prefix): # The command_completer function takes in the parser, automatically # completes optional arguments (ex, '-v'/'--verbose') or sub-commands, # and complete any arguments' values by calling a callback function # with the same arguments as complete if the callback was defined # when the parser was created. return command_completer(self.parser, shell, args, prefix) def setup(self, shell): rc = 0 if 'macros' not in shell.ctx: shell.ctx.macros = {} for name in self.base_macros: rc = self.add_macro(shell, name, self.base_macros[name]) return rc def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = 0 if ns.show: if ns.delete or ns.name: self.usage_error(shell, 'incompatible arguments: -s/--show and ', '-d/--delete' if ns.delete else 'NAME') return -1 if ns.list or ns.name: self.usage_error(shell, 'incompatible arguments: -s/--show and ', '-l/--list' if ns.list else 'NAME') return -1 if ns.show in shell.ctx.macros: print("macro ", ns.show, sep='') for line in shell.ctx.macros[ns.show]: print(" ", line, sep='') print("end") else: self.error(shell, "unknown macro ", ns.show) rc = -1 elif ns.delete: if ns.list or ns.name: self.usage_error(shell, 'incompatible arguments: -d/--delete and ', '-l/--list' if ns.list else 'NAME') return -1 if ns.delete in shell.ctx.macros: del shell.ctx.macros[ns.delete] # It gets registered as a command too. See line 230 in this # file and register() in shell.py del shell.commands[ns.delete] else: self.error(shell, "unknown macro ", ns.delete) rc = -1 elif ns.name: if ns.list: self.usage_error(shell, "list option does not take an argument") else: if (ns.name in shell.commands.keys() and ns.name not in shell.ctx.macros): self.error( shell, "A macro cannot have the same name as an ", "existing command." ) return -1 self.macro_name = ns.name self.begin_block(shell) if sys.stdin.isatty(): print("Beginning macro, use the '", shell.ctx.block.end_cmd, "' command to save.", sep='') shell.ctx.macro_orig_eof_is_sigint = shell.eof_is_sigint shell.eof_is_sigint = True elif ns.list: # Left justified table print(title_str("Registered Macros", shell.width)) chunk_size = 3 tbl = Table( columns=(Column(''), Column(''), Column('', Column.Grow)), spacing=4, header=False, width=shell.width ) macro_names = list(shell.ctx.macros.keys()) for i in range(0, len(macro_names), chunk_size): chunk = macro_names[i:i + chunk_size] tbl.extend(chunk) tbl.write(sys.stdout) else: self.usage_error(shell, "missing required argument: NAME") rc = 1 return rc def end_block(self, shell, lines): self.add_macro(shell, self.macro_name, lines) self.macro_name = None shell.eof_is_sigint = shell.ctx.macro_orig_eof_is_sigint return 0 def cancel_block(self, shell): self.macro_name = None shell.eof_is_sigint = shell.ctx.macro_orig_eof_is_sigint def add_macro(self, shell, name, lines): shell.register( Macro(lines=lines, name=name, topic='__hidden__') ) shell.ctx.macros[name] = lines return 0
class IncludeCommand(Command): ''' Execute a script file. This entails reading each line of the given file and processing it as if it were typed into the shell. ''' def __init__(self, name='include', topic='shell', brief='execute a script file', **kwargs): self.parser = PypsiArgParser( prog=name, description=brief ) self.parser.add_argument( 'path', metavar='PATH', action='store', help='file to execute' ) super(IncludeCommand, self).__init__( name=name, topic=topic, brief=brief, usage=self.parser.format_help(), **kwargs ) self.stack = [] def complete(self, shell, args, prefix): return path_completer(args[-1], prefix=prefix) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code return self.include_file(shell, ns.path) def include_file(self, shell, path): fp = None ifile = IncludeFile(path) if self.stack: for i in self.stack: if i.abspath == ifile.abspath: self.error(shell, "recursive include for file ", ifile.abspath, '\n') return -1 self.stack.append(ifile) try: fp = safe_open(path, 'r') except (OSError, IOError) as e: self.error(shell, "error opening file {}: {}".format(path, str(e))) return -1 try: # Execute the lines in the file shell.include(fp) except Exception as e: self.error(shell, "error executing file ", path, ": ", str(e)) self.stack.pop() fp.close() return 0
class MacroCommand(BlockCommand): ''' Record and execute command macros. Macros can consist of one or more command statements. Macros are comparable to Bash functions. Once a macro has been recorded, a new command is registered in the shell as an instance of a :class:`Macro`. This command requires the :class:`pypsi.plugins.block.BlockPlugin` plugin. ''' def __init__(self, name='macro', topic='shell', brief="manage registered macros", macros=None, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=MacroCmdUsage) self.parser.add_argument('-l', '--list', help='list all macros', action='store_true') self.parser.add_argument('-d', '--delete', help='delete macro', metavar='NAME', completer=self.complete_macros) self.parser.add_argument('-s', '--show', help='print macro body', metavar='NAME', completer=self.complete_macros) self.parser.add_argument('name', help='macro name', nargs='?', metavar='NAME') super(MacroCommand, self).__init__(name=name, usage=self.parser.format_help(), brief=brief, topic=topic, **kwargs) self.base_macros = macros or {} self.macro_name = None def complete_macros(self, shell, args, prefix): # pylint: disable=unused-argument # returns a list of macro names in the current shell return list(shell.ctx.macros.keys()) def complete(self, shell, args, prefix): # The command_completer function takes in the parser, automatically # completes optional arguments (ex, '-v'/'--verbose') or sub-commands, # and complete any arguments' values by calling a callback function # with the same arguments as complete if the callback was defined # when the parser was created. return command_completer(self.parser, shell, args, prefix) def setup(self, shell): rc = 0 if 'macros' not in shell.ctx: shell.ctx.macros = {} for name in self.base_macros: rc = self.add_macro(shell, name, self.base_macros[name]) return rc def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = 0 if ns.show: if ns.delete or ns.name: self.usage_error(shell, 'incompatible arguments: -s/--show and ', '-d/--delete' if ns.delete else 'NAME') return -1 if ns.list or ns.name: self.usage_error(shell, 'incompatible arguments: -s/--show and ', '-l/--list' if ns.list else 'NAME') return -1 if ns.show in shell.ctx.macros: print("macro ", ns.show, sep='') for line in shell.ctx.macros[ns.show]: print(" ", line, sep='') print("end") else: self.error(shell, "unknown macro ", ns.show) rc = -1 elif ns.delete: if ns.list or ns.name: self.usage_error(shell, 'incompatible arguments: -d/--delete and ', '-l/--list' if ns.list else 'NAME') return -1 if ns.delete in shell.ctx.macros: del shell.ctx.macros[ns.delete] # It gets registered as a command too. See line 230 in this # file and register() in shell.py del shell.commands[ns.delete] else: self.error(shell, "unknown macro ", ns.delete) rc = -1 elif ns.name: if ns.list: self.usage_error(shell, "list option does not take an argument") else: if (ns.name in shell.commands.keys() and ns.name not in shell.ctx.macros): self.error(shell, "A macro cannot have the same name as an ", "existing command.") return -1 self.macro_name = ns.name self.begin_block(shell) if sys.stdin.isatty(): print("Beginning macro, use the '", shell.ctx.block.end_cmd, "' command to save.", sep='') shell.ctx.macro_orig_eof_is_sigint = shell.eof_is_sigint shell.eof_is_sigint = True elif ns.list: # Left justified table print(title_str("Registered Macros", shell.width)) chunk_size = 3 tbl = Table(columns=(Column(''), Column(''), Column('', Column.Grow)), spacing=4, header=False, width=shell.width) macro_names = list(shell.ctx.macros.keys()) for i in range(0, len(macro_names), chunk_size): chunk = macro_names[i:i + chunk_size] tbl.extend(chunk) tbl.write(sys.stdout) else: self.usage_error(shell, "missing required argument: NAME") rc = 1 return rc def end_block(self, shell, lines): self.add_macro(shell, self.macro_name, lines) self.macro_name = None shell.eof_is_sigint = shell.ctx.macro_orig_eof_is_sigint return 0 def cancel_block(self, shell): self.macro_name = None shell.eof_is_sigint = shell.ctx.macro_orig_eof_is_sigint def add_macro(self, shell, name, lines): shell.register(Macro(lines=lines, name=name, topic='__hidden__')) shell.ctx.macros[name] = lines return 0
class HelpCommand(Command): ''' Provides access to manpage-esque topics and command usage information. ''' def __init__(self, name='help', topic='shell', brief='print information on a topic or command', topics=None, vars=None, **kwargs): self.parser = PypsiArgParser( prog=name, description=brief ) self.parser.add_argument( "topic", metavar="TOPIC", help="command or topic to print", nargs='?' ) super(HelpCommand, self).__init__( name=name, brief=brief, usage=self.parser.format_help(), topic=topic, **kwargs ) self.vars = vars or {} self.topics = topics def setup(self, shell): shell.ctx.topics = list(self.topics or []) shell.ctx.uncat_topic = Topic('uncat', 'Uncategorized Commands & Topics') shell.ctx.topic_lookup = {t.id: t for t in shell.ctx.topics} shell.ctx.topics_dirty = True def complete(self, shell, args, prefix): if shell.ctx.topics_dirty: self.reload(shell) completions = [ x.id for x in shell.ctx.topics if x.id.startswith(prefix) or not prefix ] completions.extend([ x for x in shell.commands if x.startswith(prefix) or not prefix ]) completions = sorted(completions) return completions def reload(self, shell): shell.ctx.uncat_topic.commands = [] for id in shell.ctx.topic_lookup: shell.ctx.topic_lookup[id].commands = [] for (name, cmd) in shell.commands.items(): if cmd.topic == '__hidden__': continue if cmd.topic: if cmd.topic in shell.ctx.topic_lookup: shell.ctx.topic_lookup[cmd.topic].commands.append(cmd) else: self.add_topic(shell, Topic(cmd.topic, commands=[cmd])) else: shell.ctx.uncat_topic.commands.append(cmd) shell.ctx.topics_dirty = False for topic in shell.ctx.topics: if topic.commands: topic.commands = sorted(topic.commands, key=lambda x: x.name) def add_topic(self, shell, topic): shell.ctx.topics_dirty = True shell.ctx.topic_lookup[topic.id] = topic shell.ctx.topics.append(topic) def print_topic_commands(self, shell, topic, title=None, name_col_width=20): print( AnsiCodes.yellow, title_str(title or topic.name or topic.id, shell.width), AnsiCodes.reset, sep='' ) print(AnsiCodes.yellow, end='') Table( columns=(Column(''), Column('', Column.Grow)), spacing=4, header=False, width=shell.width ).extend(*[ (' ' + c.name.ljust(name_col_width - 1), c.brief or '') for c in topic.commands ]).write(sys.stdout) print(AnsiCodes.reset, end='') def print_topics(self, shell): max_name_width = 0 for topic in shell.ctx.topics: for c in topic.commands: max_name_width = max(len(c.name), max_name_width) for c in shell.ctx.uncat_topic.commands: max_name_width = max(len(c.name), max_name_width) addl = [] for topic in shell.ctx.topics: if topic.content or not topic.commands: addl.append(topic) if topic.commands: self.print_topic_commands(shell, topic, name_col_width=max_name_width) print() if shell.ctx.uncat_topic.commands: self.print_topic_commands(shell, shell.ctx.uncat_topic, name_col_width=max_name_width) print() if addl: addl = sorted(addl, key=lambda x: x.id) print( AnsiCodes.yellow, title_str("Additional Topics", shell.width), sep='' ) Table( columns=(Column(''), Column('', Column.Grow)), spacing=4, header=False, width=shell.width ).extend(*[ (' ' + topic.id.ljust(max_name_width - 1), topic.name or '') for topic in addl ]).write(sys.stdout) print(AnsiCodes.reset) def print_topic(self, shell, id): if id not in shell.ctx.topic_lookup: if id in shell.commands: cmd = shell.commands[id] print(AnsiCodes.yellow, cmd.usage, AnsiCodes.reset, sep='') return 0 self.error(shell, "unknown topic: ", id) return -1 topic = shell.ctx.topic_lookup[id] if topic.content: print(title_str(topic.name or topic.id, shell.width)) try: cnt = topic.content.format(**self.vars) except: cnt = topic.content print(cnt) print() if topic.commands: self.print_topic_commands(shell, topic, "Commands") return 0 def run(self, shell, args): if shell.ctx.topics_dirty: self.reload(shell) try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = 0 if not ns.topic: self.print_topics(shell) else: rc = self.print_topic(shell, ns.topic) return rc
class InputOutputCommand(Command): SUBCOMMAND_DELETE = "delete" SUBCOMMAND_EDIT = "edit" SUBCOMMAND_ADD = "add" def __init__(self, name, topic, brief, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief) # TODO: start using args from ArgParser, now they only added to show pretty help self.parser = PypsiArgParser(prog=name, description=brief) subparsers = self.parser.add_subparsers(dest='subparser_name', help="Commands") add_arg_parser = subparsers.add_parser("add", help="add connection") edit_arg_parser = subparsers.add_parser("edit", help="edit connection") edit_arg_parser.add_argument("index", metavar="INDEX", action="store", help="connection index to edit") delete_arg_parser = subparsers.add_parser("delete", help="delete connection") delete_arg_parser.add_argument("index", metavar="INDEX", action="store", help="connection index to delete") super(InputOutputCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) self._handlers = { self.SUBCOMMAND_ADD: self._add_handler, self.SUBCOMMAND_DELETE: self._delete_handler, self.SUBCOMMAND_EDIT: self._edit_handler } self._print_header = True def run(self, shell, args): try: self.parser.parse_args(args) except CommandShortCircuit as e: return e.code component_info = shell._component_info self._print_header = shell._wizard_edit_mode if self.name == CommandName.INPUT: if component_info._inputs is None: component_info._inputs = [] self._inputs_or_outputs_array = component_info._inputs elif self.name == CommandName.OUTPUT: if component_info._outputs is None: component_info._outputs = [] self._inputs_or_outputs_array = component_info._outputs cmd = self.SUBCOMMAND_ADD if args is not None and len(args): if args[0] not in self._handlers.keys(): self.usage_error(shell, args) return 1 cmd = args[0] return self._handlers[cmd](shell, args) def _add_handler(self, shell, args=None): comp_conn = ComponentConnectionInfo() wizard_header = "{} Connection Configuration".format( self.name.capitalize()) ns = ConnectionWizard(wizard_header, comp_conn).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_conn.load_from_json(out_dict) self._inputs_or_outputs_array.append(comp_conn) return 0 def _edit_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 try: index = int(args[1]) except: print("Index parameter '{}' can not be casted to int".format( args[1])) self.usage_error(shell, args) return 0 if 0 <= index < len(self._inputs_or_outputs_array): comp_conn = self._inputs_or_outputs_array[index] wizard_header = "{} Connection Configuration".format( self.name.capitalize()) ns = ConnectionWizard(wizard_header, comp_conn).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_conn.load_from_json(out_dict) else: print("Element with index '{}' is not found".format(index)) self.usage_error(shell, args) return 0 def _delete_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 try: index = int(args[1]) except: print("Index parameter '{}' can not be casted to int".format( args[1])) self.usage_error(shell, args) return 0 if 0 <= index < len(self._inputs_or_outputs_array): self._inputs_or_outputs_array.pop(index) else: print("Element with index '{}' is not found".format(index)) self.usage_error(shell, args) return 0
class TipCommand(Command): def __init__(self, name='tip', brief='print shell tips', topic='shell', tips=None, motd=None, vars=None, **kwargs): self.tips = tips or [] self.motd = motd or '' self.vars = vars or {} self.rand = random.Random() self.rand.seed() self.parser = PypsiArgParser( prog=name, description=brief ) self.parser.add_argument( '-m', '--motd', action='store_true', help="print the message of the day" ) super(TipCommand, self).__init__( name=name, brief=brief, topic=topic, usage=self.parser.format_help(), **kwargs ) def load_tips(self, path): fp = safe_open(path, 'r') tip = [] for line in fp.readlines(): line = line.rstrip() if line: tip.append(line) elif tip: self.tips.append(' '.join(tip)) tip = [] if tip: self.tips.append(' '.join(tip)) fp.close() def load_motd(self, path): fp = safe_open(path, 'r') self.motd = fp.read().rstrip() fp.close() def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = None if ns.motd: rc = self.print_motd(shell) else: rc = self.print_random_tip(shell) return rc def print_random_tip(self, shell, header=True): if not self.tips: self.error(shell, "no tips available") return -1 i = self.rand.randrange(len(self.tips)) if header: title = "Tip #{}\n".format(i + 1) title += '-' * len(title) print(AnsiCodes.green, title, AnsiCodes.reset, sep='') try: cnt = sys.stdout.ansi_format(self.tips[i], **self.vars) except: cnt = self.tips[i] print(cnt) return 0 def print_motd(self, shell): if not self.motd: self.error(shell, "no motd available") return -1 try: cnt = sys.stdout.ansi_format(self.motd, **self.vars) except: cnt = self.motd print( AnsiCodes.green, "Message of the Day".center(shell.width), '\n', AnsiCodes.green, '>' * shell.width, "\n", AnsiCodes.reset, cnt, '\n', AnsiCodes.green, "<" * shell.width, "\n", AnsiCodes.reset, sep='' ) return 0
class IncludeCommand(Command): ''' Execute a script file. This entails reading each line of the given file and processing it as if it were typed into the shell. ''' def __init__(self, name='include', topic='shell', brief='execute a script file', **kwargs): self.parser = PypsiArgParser(prog=name, description=brief) self.parser.add_argument('path', metavar='PATH', action='store', help='file to execute') super(IncludeCommand, self).__init__(name=name, topic=topic, brief=brief, usage=self.parser.format_help(), **kwargs) self.stack = [] def complete(self, shell, args, prefix): return path_completer(args[-1]) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code return self.include_file(shell, ns.path) def include_file(self, shell, path): fp = None ifile = IncludeFile(path) if self.stack: for i in self.stack: if i.abspath == ifile.abspath: self.error(shell, "recursive include for file ", ifile.abspath, '\n') return -1 self.stack.append(ifile) try: fp = safe_open(path, 'r') except (OSError, IOError) as e: self.error(shell, "error opening file ", path, ": ", str(e), '\n') return -1 for line in fp: shell.execute(line.strip()) ifile.line += 1 self.stack.pop() fp.close() return 0
class ArgumentCommand(Command): SUBCOMMAND_DELETE = "delete" SUBCOMMAND_EDIT = "edit" SUBCOMMAND_ADD = "add" def __init__(self, name, topic, brief, **kwargs): # TODO: start using args from ArgParser, now they only added to show pretty help self.parser = PypsiArgParser(prog=name, description=brief) subparsers = self.parser.add_subparsers(dest='subparser_name', help="Commands") add_arg_parser = subparsers.add_parser("add", help="add argument") edit_arg_parser = subparsers.add_parser("edit", help="edit argument") edit_arg_parser.add_argument("key", metavar="KEY", action="store", help="argument key to edit") delete_arg_parser = subparsers.add_parser("delete", help="delete argument") delete_arg_parser.add_argument("key", metavar="KEY", action="store", help="argument key to delete") super(ArgumentCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) self._handlers = { self.SUBCOMMAND_ADD: self._add_handler, self.SUBCOMMAND_DELETE: self._delete_handler, self.SUBCOMMAND_EDIT: self._edit_handler } def run(self, shell, args): try: self.parser.parse_args(args) except CommandShortCircuit as e: return e.code cmd = self.SUBCOMMAND_ADD self._print_header = shell._wizard_edit_mode if args is not None and len(args): if args[0] not in self._handlers.keys(): self.usage_error(shell, args) return 1 cmd = args[0] return self._handlers[cmd](shell, args) def _add_handler(self, shell, args=None): component_info = shell._component_info comp_arg = ComponentArgumentInfo() ns = ArgumentWizard(component_info, comp_arg).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_arg.load_from_json(out_dict) if component_info.arguments is None: component_info.arguments = [] component_info.arguments.append(comp_arg) return 0 def _edit_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 component_info = shell._component_info key = args[1] comp_arg = component_info.get_argument(key) if comp_arg: ns = ArgumentWizard(component_info, comp_arg).run(shell, self._print_header) if ns: out_dict = Helper.filter_non_zero_string_items(ns.__dict__) comp_arg.load_from_json(out_dict) else: print("Argument with the key '{}' is not found".format(key)) self.usage_error(shell, args) return 0 def _delete_handler(self, shell, args): if len(args) < 2: self.usage_error(shell, args) return 1 component_info = shell._component_info key = args[1] arg = component_info.get_argument(key) if arg: component_info.arguments.remove(arg) else: print("Argument with the key '{}' is not found".format(key)) self.usage_error(shell, args) return 0
class MinionCommand(Command): def __init__(self, topic='proxy', name='minion', brief=HelpTopic, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=MinionCmdUsage) subcmd = self.parser.add_subparsers(prog='minion', dest='subcmd', metavar='subcmd') subcmd.required = True shutdown = subcmd.add_parser('shutdown', help='shutdown the connected minion') restart = subcmd.add_parser('restart', help='restart the connected minion') load = subcmd.add_parser('load', help='loads a plugin into minion') load.add_argument('plugin', metavar='PLUGIN', help='name of plugin to load into minion') log = subcmd.add_parser('log', help='print the log of a minion') log.add_argument('lines', metavar='LINES', help='range of line number to print <start:finish>') super(MinionCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): if not self.check_connected(shell): self.error("shell", "Not currently connected to minion") return 0 try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if ns.subcmd == 'shutdown': self.shutdown_minion(shell) elif ns.subcmd == 'restart': self.restart_minion(shell) elif ns.subcmd == 'load': self.load_plugin(shell, ns.plugin) def complete(self, shell, args, prefix): cmds = ['shutdown', 'restart', 'load', 'log'] if len(args) == 1: return sorted([cmd for cmd in cmds if cmd.startswith(prefix)]) def check_connected(self, shell): minion = shell.connect if minion is not None: return True else: False def shutdown_minion(self, shell): params = minionRPC_pb2.Empty() rpc = shell.minions[shell.connect]['rpc'] try: ret = rpc.shutdown(params) except Exception as e: shell.log.error(str(e)) self.error('shell', f"Failed to shutdown minion: {shell.connect}") if ret.reply == "okay": print(f'Shutting down minion: {shell.connect}') else: self.error('shell', f"Failed to shutdown minion: {shell.connect}") def restart_minion(self, shell): params = minionRPC_pb2.Empty() rpc = shell.minions[shell.connect]['rpc'] try: ret = rpc.runPluginFunction(params) except Exception as e: shell.log.error(str(e)) self.error('shell', f"Failed to restart minion: {shell.connect}") if ret.reply == "okay": print(f'Restarting minion: {shell.connect}') else: self.error('shell', f"Failed to restart minion: {shell.connect}") def load_plugin(self, shell, plugin): try: plg = LoaderUtil.load_plugin(plugin) except: self.error( 'shell', f"Plugin: {plugin} is not in the available plugin list") params = minionRPC_pb2.loadReq(plugin=plugin) rpc = shell.minions[shell.connect]['rpc'] try: ret = rpc.loadPlugin(params) except Exception as e: shell.log.error(str(e)) if ret.error < 1: print(f'Loaded plugin: {plugin} into minion: {shell.connect}') else: shell.log.error(ret.reply) self.error( 'shell', f"Failed to load plugin: { plugin } in {shell.connect}")
class TailCommand(Command): ''' Displays the last N lines of a file to the screen. ''' def __init__(self, name='tail', topic='shell', brief='display the last lines of a file', **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=TailCmdUsage) # Add a callback function that will be called when the # argument is tab-completed self.parser.add_argument('input_file', help='file to display', metavar="FILE", completer=self.complete_path) self.parser.add_argument('-n', '--lines', metavar="N", type=int, default=10, help="number of lines to display") self.parser.add_argument('-f', '--follow', help="continue to output as file grows", action='store_true') super(TailCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code # check for valid input file if not os.path.isfile(ns.input_file): self.error(shell, "invalid file path: ", ns.input_file, "\n") return -1 # print the last N lines last_lines = self.tail(ns.input_file, ns.lines) for line in last_lines: print(line) print() # continue to follow the file and display new content if ns.follow: self.follow_file(ns.input_file) return 0 def complete_path(self, shell, args, prefix): # pylint: disable=unused-argument return path_completer(args[-1]) def complete(self, shell, args, prefix): # The command_completer function takes in the parser, automatically # completes optional arguments (ex, '-v'/'--verbose') or sub-commands, # and complete any arguments' values by calling a callback function # with the same arguments as complete if the callback was defined # when the parser was created. return command_completer(self.parser, shell, args, prefix) def tail(self, fname, lines=10, block_size=1024): data = [] blocks = -1 num_lines = 0 with open(fname) as fp: # seek to the end and get the number of bytes in the file fp.seek(0, 2) num_bytes = fp.tell() bytes_left = num_bytes while num_lines < lines and bytes_left > 0: if bytes_left - block_size > 0: # Seek back a block_size fp.seek(num_bytes - (blocks * block_size)) # read data from file data.insert(0, fp.read(block_size)) else: # jump back to the beginning fp.seek(0, 0) # read data data.insert(0, fp.read(num_bytes)) num_lines = data[0].count('\n') bytes_left -= block_size blocks -= 1 return ''.join(data).splitlines()[-lines:] def follow_file(self, fname): with open(fname) as fp: # jump to the end of the file fp.seek(0, 2) try: while 1: where = fp.tell() line = fp.readline() if not line: time.sleep(1) fp.seek(where) else: print(line, end='', flush=True) except KeyboardInterrupt: print() return 0 return 0
class VariableCommand(Command): ''' Manage variables. ''' Usage = """var name or: var EXPRESSION or: var -l or: var -d name""" def __init__(self, name='var', brief='manage local variables', topic='shell', **kwargs): self.parser = PypsiArgParser( prog=name, description=brief, usage=VariableCommand.Usage ) self.parser.add_argument( '-l', '--list', help='list variables', action='store_true' ) self.parser.add_argument( '-d', '--delete', help='delete variable', metavar='VARIABLE' ) self.parser.add_argument( 'exp', metavar='EXPRESSION', help='expression defining variable, in the form of NAME = [VALUE]', nargs=argparse.REMAINDER ) super(VariableCommand, self).__init__( name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs ) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code rc = 0 if ns.list: tbl = Table( columns=(Column("Variable"), Column("Value", Column.Grow)), width=shell.width, spacing=4, ) for name in shell.ctx.vars: if name == '_': continue s = shell.ctx.vars[name] if callable(s): s = s() elif isinstance(s, ManagedVariable): s = s.get(shell) tbl.append(name, obj_str(s)) tbl.write(sys.stdout) elif ns.delete: if ns.delete in shell.ctx.vars: s = shell.ctx.vars[ns.delete] if isinstance(s, ManagedVariable): self.error(shell, "variable is managed and cannot be deleted") rc = -1 else: del shell.ctx.vars[ns.delete] else: self.error(shell, "unknown variable: ", ns.delete) elif ns.exp and '=' in ''.join(args): (remainder, exp) = Expression.parse(args) if remainder or not exp: self.error(shell, "invalid expression") return 1 if isinstance(shell.ctx.vars[exp.operand], ManagedVariable): try: shell.ctx.vars[exp.operand].set(shell, exp.value) except ValueError as e: self.error(shell, "could not set variable: ", str(e)) else: shell.ctx.vars[exp.operand] = exp.value elif ns.exp: if len(args) == 1: if args[0] in shell.ctx.vars: s = shell.ctx.vars[args[0]] if callable(s): s = s() elif isinstance(s, ManagedVariable): s = s.getter(shell) print(obj_str(s)) else: self.error(shell, "unknown variable: ", args[0]) return 1 else: self.error(shell, "invalid expression") return 1 else: self.usage_error(shell, "missing required EXPRESSION") rc = 1 return rc
class LogCommand(Command): def __init__(self, topic='proxy', name='log', brief=HelpTopic, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=LogCmdUsage) self.parser.add_argument('-a', '--all', help="prints out the entire log file", action='store_true') self.parser.add_argument( '-l', '--lines', help='prints out a range of lines <start_line>:<end_line>', metavar="LINES") super(LogCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if ns.all and ns.lines: self.error( 'shell', "Cannot specify -l/--lines and -a/--all at the sametime") return 1 elif ns.all: self.print_logfile(shell, 0, 0, True) elif ns.lines: lnum = ns.lines.split(":") if len(lnum) > 2 or len(lnum) < 2: self.error( 'shell', "improper format: expecting <start_line>:<end_line>") return 1 else: self.print_logfile(shell, lnum[0], lnum[1], False) def print_logfile(self, shell, start, end, printall): path = shell.config.log.get('path', 'master.log') red = AnsiCodes.red r = AnsiCodes.reset yellow = AnsiCodes.yellow if path == 'master.log': path = OsUtil.get_master_log_path() try: fh = open(path, "r") count = 0 if start == '': start = 0 start = int(start) if end != '': end = int(end) for line in fh: if count >= start or printall: if 'ERROR' in line or 'EXCEPTION' in line: sys.stdout.write(f"{red}{line}{r}") elif 'WARN' in line: sys.stdout.write(f"{yellow}{line}{r}") else: sys.stdout.write(line) if (end != '') and ((count + 1) > end) and not printall: fh.close() return 0 count += 1 fh.close() except Exception as e: self.error('shell', f"could not display logfile {path}")
class ConnectCommand(Command): def __init__(self, topic='proxy', name='connect', brief=HelpTopic, **kwargs): self.parser = PypsiArgParser(prog=name, description=brief, usage=ConnectCmdUsage) self.parser.add_argument('name', help='name of minion', metavar="NAME", completer=self.complete_minions) super(ConnectCommand, self).__init__(name=name, usage=self.parser.format_help(), topic=topic, brief=brief, **kwargs) def run(self, shell, args): try: ns = self.parser.parse_args(args) except CommandShortCircuit as e: return e.code if shell.connect is not None: self.error( shell, f"could not connect to {ns.name}: currently connected to {shell.connect}" ) elif ns.name in shell.minions: if self.test_minion_connection(shell, ns.name): shell.connect = ns.name shell.add_prompt_connect() print(f"Successfully connected to minion: {ns.name}") else: self.error( shell, f"could not connect to {ns.name}: failed alive check") else: self.error( shell, f"cannot connect to {ns.name}: not a configured minion") def complete(self, shell, args, prefix): if len(args) == 1: return self.complete_minions(shell, args, prefix) def complete_minions(self, shell, args, prefix): return sorted( [name for name in shell.minions if name.startswith(prefix)]) def test_minion_connection(self, shell, minion): try: rpc = shell.minions[minion]['rpc'] msg = minionRPC_pb2.Empty() ret = rpc.checkAlive(msg) if ret.reply == 'alive': return True else: return False except Exception as e: shell.log.error(str(e)) return False