From 9ddaf96ccf53dbd7264c33949a58ca82735bc179 Mon Sep 17 00:00:00 2001 From: n1nj4sec Date: Sat, 10 Oct 2015 15:41:02 +0200 Subject: [PATCH] PupyArgumentParser.add_argument can now take a "completer" keywork argument to change the command line autocompletion for a specific argument ! useful to use with path_completer --- .gitignore | 1 + pupy/modules/download.py | 3 +- pupy/modules/memory_exec.py | 5 +- pupy/modules/persistence.py | 17 ++- pupy/modules/upload.py | 3 +- pupy/pupylib/PupyCmd.py | 44 ++++++-- pupy/pupylib/PupyCompleter.py | 182 ++++++++++++++++++++++++++++++++ pupy/pupylib/PupyModule.py | 25 ++++- pupy/pupylib/PupyServer.py | 6 ++ pupy/pupylib/PythonCompleter.py | 6 +- 10 files changed, 277 insertions(+), 15 deletions(-) create mode 100644 pupy/pupylib/PupyCompleter.py diff --git a/.gitignore b/.gitignore index 50f198d7..31ab5a5d 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ pupy/data/ pupy/.pupy_history .DS_Store +*.swp #do not redistribute microsoft visual C++ DLLs (LICENSE) diff --git a/pupy/modules/download.py b/pupy/modules/download.py index 74b55e45..948f81ff 100644 --- a/pupy/modules/download.py +++ b/pupy/modules/download.py @@ -1,5 +1,6 @@ # -*- coding: UTF8 -*- from pupylib.PupyModule import * +from pupylib.PupyCompleter import * from rpyc.utils.classic import download import os import os.path @@ -11,7 +12,7 @@ class DownloaderScript(PupyModule): def init_argparse(self): self.arg_parser = PupyArgumentParser(prog='download', description=self.__doc__) self.arg_parser.add_argument('remote_file', metavar='') - self.arg_parser.add_argument('local_file', nargs='?', metavar='') + self.arg_parser.add_argument('local_file', nargs='?', metavar='', completer=path_completer) def run(self, args): remote_file=self.client.conn.modules['os.path'].expandvars(args.remote_file) rep=os.path.join("data","downloads",self.client.short_name()) diff --git a/pupy/modules/memory_exec.py b/pupy/modules/memory_exec.py index e3a2242d..49dc519a 100644 --- a/pupy/modules/memory_exec.py +++ b/pupy/modules/memory_exec.py @@ -14,6 +14,7 @@ # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE # -------------------------------------------------------------- from pupylib.PupyModule import * +from pupylib.PupyCompleter import * __class_name__="MemoryExec" @@ -25,9 +26,9 @@ class MemoryExec(PupyModule): self.interrupted=False self.mp=None def init_argparse(self): - self.arg_parser = PupyArgumentParser(prog="msgbox", description=self.__doc__) + self.arg_parser = PupyArgumentParser(prog="memory_exec", description=self.__doc__) self.arg_parser.add_argument('-p', '--process', default='cmd.exe', help='process to start suspended') - self.arg_parser.add_argument('--fork', action='store_true', help='fork and do not wait for the child program. stdout will not be retrieved') + self.arg_parser.add_argument('--fork', action='store_true', help='fork and do not wait for the child program. stdout will not be retrieved', completer=path_completer) self.arg_parser.add_argument('--interactive', action='store_true', help='interactive with the new process stdin/stdout') self.arg_parser.add_argument('path', help='path to the exe') self.arg_parser.add_argument('args', nargs='*', help='optional arguments to pass to the exe') diff --git a/pupy/modules/persistence.py b/pupy/modules/persistence.py index 9d9cb7c8..1d8c2169 100644 --- a/pupy/modules/persistence.py +++ b/pupy/modules/persistence.py @@ -1,5 +1,20 @@ # -*- coding: UTF8 -*- +# -------------------------------------------------------------- +# Copyright (c) 2015, Nicolas VERDIER (contact@n1nj4.eu) +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE +# -------------------------------------------------------------- from pupylib.PupyModule import * +from pupylib.PupyCompleter import * import random import pupygen import os.path @@ -11,7 +26,7 @@ class PersistenceModule(PupyModule): """ Enables persistence via registry keys """ def init_argparse(self): self.arg_parser = PupyArgumentParser(prog="persistence", description=self.__doc__) - self.arg_parser.add_argument('-e','--exe', help='Use an alternative file and set persistency') + self.arg_parser.add_argument('-e','--exe', help='Use an alternative file and set persistency', completer=path_completer) self.arg_parser.add_argument('-m','--method', choices=['registry'], required=True, help='persistence method') @windows_only diff --git a/pupy/modules/upload.py b/pupy/modules/upload.py index 67c6de66..6e585237 100644 --- a/pupy/modules/upload.py +++ b/pupy/modules/upload.py @@ -1,5 +1,6 @@ # -*- coding: UTF8 -*- from pupylib.PupyModule import * +from pupylib.PupyCompleter import * from rpyc.utils.classic import upload import os import os.path @@ -10,7 +11,7 @@ class UploaderScript(PupyModule): """ upload a file/directory to a remote system """ def init_argparse(self): self.arg_parser = PupyArgumentParser(prog='download', description=self.__doc__) - self.arg_parser.add_argument('local_file', metavar='') + self.arg_parser.add_argument('local_file', metavar='', completer=path_completer) self.arg_parser.add_argument('remote_file', metavar='') def run(self, args): upload(self.client.conn, args.local_file, args.remote_file) diff --git a/pupy/pupylib/PupyCmd.py b/pupy/pupylib/PupyCmd.py index 9acbdfa9..76b7fff9 100644 --- a/pupy/pupylib/PupyCmd.py +++ b/pupy/pupylib/PupyCmd.py @@ -1,3 +1,4 @@ +# -*- coding: UTF8 -*- # -------------------------------------------------------------- # Copyright (c) 2015, Nicolas VERDIER (contact@n1nj4.eu) # All rights reserved. @@ -38,10 +39,11 @@ import logging import traceback import rpyc import rpyc.utils.classic -from .PythonCompleter import PupyCompleter +from .PythonCompleter import PythonCompleter from .PupyErrors import PupyModuleExit, PupyModuleError from .PupyModule import PupyArgumentParser from .PupyJob import PupyJob +from .PupyCompleter import PupyCompleter import argparse from pupysh import __version__ import copy @@ -169,7 +171,6 @@ class PupyCmd(cmd.Cmd): self.intro = color(BANNER, 'green') self.prompt = color('>> ','blue', prompt=True) self.doc_header = 'Available commands :\n' - self.complete_space=['run'] self.default_filter=None try: if not self.config.getboolean("cmdline","display_banner"): @@ -183,6 +184,7 @@ class PupyCmd(cmd.Cmd): self.aliases[command]=alias except Exception as e: logging.warning("error while parsing aliases from pupy.conf ! %s"%str(traceback.format_exc())) + self.pupy_completer=PupyCompleter(self.aliases, self.pupsrv) @staticmethod def table_format(diclist, wl=[], bl=[]): @@ -270,9 +272,7 @@ class PupyCmd(cmd.Cmd): def completenames(self, text, *ignored): dotext = 'do_'+text - if text in self.complete_space: - return [a[3:]+" " for a in self.get_names() if a.startswith(dotext)]+[x+" " for x in self.aliases.iterkeys() if x.startswith(text)] - return [a[3:] for a in self.get_names() if a.startswith(dotext)]+[x for x in self.aliases.iterkeys() if x.startswith(text)] + return [a[3:]+' ' for a in self.get_names() if a.startswith(dotext)]+[x+' ' for x in self.aliases.iterkeys() if x.startswith(text)] def pre_input_hook(self): #readline.redisplay() @@ -495,7 +495,7 @@ class PupyCmd(cmd.Cmd): oldcompleter=readline.get_completer() try: local_ns={"pupsrv":self.pupsrv} - readline.set_completer(PupyCompleter(local_ns=local_ns).complete) + readline.set_completer(PythonCompleter(local_ns=local_ns).complete) readline.parse_and_bind('tab: complete') code.interact(local=local_ns) except Exception as e: @@ -596,9 +596,41 @@ class PupyCmd(cmd.Cmd): if pj: del pj + def complete(self, text, state): + if state == 0: + import readline + origline = readline.get_line_buffer() + line = origline.lstrip() + stripped = len(origline) - len(line) + begidx = readline.get_begidx() - stripped + endidx = readline.get_endidx() - stripped + if begidx>0: + cmd, args, foo = self.parseline(line) + if cmd == '': + compfunc = self.completedefault + else: + try: + #compfunc = getattr(self, 'complete_' + cmd) + compfunc = self.pupy_completer.complete + except AttributeError: + compfunc = self.completedefault + else: + compfunc = self.completenames + self.completion_matches = compfunc(text, line, begidx, endidx) + try: + return self.completion_matches[state] + except IndexError: + return None + #text : word match #line : complete line def complete_run(self, text, line, begidx, endidx): + pmc=PupyModCompleter(None) + try: + res=pmc.complete(text, line, begidx, endidx) + except Exception as e: + print e + return res mline = line.partition(' ')[2] joker=1 diff --git a/pupy/pupylib/PupyCompleter.py b/pupy/pupylib/PupyCompleter.py new file mode 100644 index 00000000..52b4f5d2 --- /dev/null +++ b/pupy/pupylib/PupyCompleter.py @@ -0,0 +1,182 @@ +# -*- coding: UTF8 -*- +# -------------------------------------------------------------- +# Copyright (c) 2015, Nicolas VERDIER (contact@n1nj4.eu) +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE +# -------------------------------------------------------------- + +import os +import os.path +import shlex +import re + +def debug(msg): + #with open("/tmp/debug.log","a+") as log: + # log.write(str(msg)+"\n") + pass + +def list_completer(l): + def func(text, line, begidx, endidx): + return [x for x in l if text.startswith(x)] + +def void_completer(text, line, begidx, endidx): + return [] + +def path_completer(text, line, begidx, endidx): + l=[] + if not text: + l=os.listdir(".") + else: + try: + dirname=os.path.dirname(text) + basename=os.path.basename(text) + for f in os.listdir(dirname): + if f.startswith(basename): + if os.path.isdir(os.path.join(dirname,f)): + l.append(os.path.join(dirname,f)+os.sep) + else: + l.append(os.path.join(dirname,f)+" ") + except Exception as e: + pass + return l + +class PupyCompleter(object): + def __init__(self, aliases, pupysrv): + self.aliases=aliases + self.pupysrv=pupysrv + + def get_module_completer(self, name): + #TODO handle aliases + return self.pupysrv.get_module_completer(name) + + def complete(self, text, line, begidx, endidx): + try: + #debug("\"%s\" \"%s\" %s %s"%(text, line, begidx, endidx)) + if line.startswith("run "): + res=self.complete_run(text, line, begidx, endidx) + if res is not None: + return res + modname=line[4:].split()[0] + completer_func=self.get_module_completer(modname).complete + debug("%s"%completer_func) + if completer_func: + return completer_func(text, line, begidx, endidx) + else: + return [] + except Exception as e: + #print e + pass + + def complete_run(self, text, line, begidx, endidx): + mline = line.partition(' ')[2] + joker=1 + found_module=False + #handle autocompletion of modules with --filter argument + for x in shlex.split(mline): + if x in ("-f", "--filter"):#arguments with a param + joker+=1 + elif x in ("--bg",):#arguments without parameter + pass + else: + joker-=1 + if not x.startswith("-") and joker==0: + found_module=True + if joker<0: + return + if ((len(text)>0 and joker==0) or (len(text)==0 and not found_module and joker<=1)): + return [re.sub(r"(.*)\.pyc?$",r"\1",x)+" " for x in os.listdir("modules") if x.startswith(text) and not x=="__init__.py" and not x=="__init__.pyc"] + + +class PupyModCompleter(object): + def __init__(self): + self.conf= { + "positional_args":[ + #TODO + ], + "optional_args":[ + ], + } + + def add_positional_arg(self, names, **kwargs): + """ names can be a string or a list to pass args aliases at once """ + if not type(names) is list and not type(names) is tuple: + names=[names] + for name in names: + self.conf["positional_args"].append((name, kwargs)) + + def add_optional_arg(self, names, **kwargs): + """ names can be a string or a list to pass args aliases at once """ + if not type(names) is list and not type(names) is tuple: + names=[names] + for name in names: + self.conf["optional_args"].append((name, kwargs)) + + def get_optional_nargs(self, name): + if "action" in self.conf["optional_args"]: + action=self.conf["optional_args"]["action"] + if action=="store_true" or action=="store_false": + return 0 + return 1 + + def get_optional_args(self, nargs=None): + if nargs is None: + return [x[0] for x in self.conf["optional_args"]] + else: + return [x[0] for x in self.conf["optional_args"] if self.get_optional_nargs(x[0])==nargs] + + def get_last_text(self, text, line, begidx, endidx): + try: + return line[0:begidx-1].rsplit(' ',1)[1].strip() + except Exception: + return None + + def get_positional_arg_index(self, text, line, begidx, endidx): + #TODO + tab=shlex.split(line) + positional_index=-1 + for i in range(0, len(tab)): + if tab[i] in self.get_optional_args(nargs=0): + continue + elif tab[i] in self.get_optional_args(nargs=1): + i+=1 + continue + else: + positional_index+=1 + if len(text)==0: + positional_index+=1 + return positional_index + + def get_optional_args_completer(self, name): + return [x[1]["completer"] for x in self.conf["optional_args"] if x[0]==name][0] + def get_positional_args_completer(self, index): + return self.conf["positional_args"][index][1]["completer"] + + def complete(self, text, line, begidx, endidx): + debug("\"%s\" \"%s\" %s %s"%(text, line, begidx, endidx)) + last_text=self.get_last_text(text, line, begidx, endidx) + debug("last text: %s"%last_text) + if last_text in self.get_optional_args(nargs=1): + debug(self.get_optional_args_completer(last_text)) + return self.get_optional_args_completer(last_text)(text, line, begidx, endidx) + if text.startswith("-"): #positional args completer + return [x+" " for x in self.get_optional_args() if x.startswith(text)] + else: + try: + debug("here") + positional_index=self.get_positional_arg_index(text, line, begidx, endidx)-2 # -2 for "run" + "module_name" + debug("positional index is %s"%positional_index) + return self.get_positional_args_completer(positional_index)(text, line, begidx, endidx) + except Exception as e: + debug(e) + + + diff --git a/pupy/pupylib/PupyModule.py b/pupy/pupylib/PupyModule.py index 1def5e7f..4d863861 100644 --- a/pupy/pupylib/PupyModule.py +++ b/pupy/pupylib/PupyModule.py @@ -16,6 +16,7 @@ import argparse import sys from .PupyErrors import PupyModuleExit +from .PupyCompleter import PupyModCompleter, void_completer import StringIO class PupyArgumentParser(argparse.ArgumentParser): @@ -24,6 +25,28 @@ class PupyArgumentParser(argparse.ArgumentParser): self._print_message(message, sys.stderr) raise PupyModuleExit("exit with status %s"%status) + def add_argument(self, *args, **kwargs): + completer_func=None + if "completer" in kwargs: + completer_func=kwargs["completer"] + del kwargs["completer"] + else: + completer_func=void_completer + argparse.ArgumentParser.add_argument(self, *args, **kwargs) + completer=self.get_completer() + for a in args: + if a.startswith("-"): + completer.add_optional_arg(a, completer=completer_func) + else: + completer.add_positional_arg(a, completer=completer_func) + + def get_completer(self): + if hasattr(self,'pupy_mod_completer') and self.pupy_mod_completer is not None: + return self.pupy_mod_completer + else: + self.pupy_mod_completer=PupyModCompleter() + return self.pupy_mod_completer + class PupyModule(object): """ This is the class all the pupy scripts must inherit from @@ -38,7 +61,6 @@ class PupyModule(object): """ client must be a PupyClient instance """ self.client=client self.job=job - self.init_argparse() if formatter is None: from .PupyCmd import PupyCmd self.formatter=PupyCmd @@ -50,6 +72,7 @@ class PupyModule(object): else: self.stdout=stdout self.del_close=False + self.init_argparse() def __del__(self): if self.del_close: diff --git a/pupy/pupylib/PupyServer.py b/pupy/pupylib/PupyServer.py index 4ea14137..3db954d3 100644 --- a/pupy/pupylib/PupyServer.py +++ b/pupy/pupylib/PupyServer.py @@ -204,6 +204,12 @@ class PupyServer(threading.Thread): l.append((module_name,module.__doc__)) return l + def get_module_completer(self, module_name): + """ return the module PupyCompleter if any is defined""" + module=self.get_module(module_name) + ps=module(None,None) + return ps.arg_parser.get_completer() + def get_module(self, name): script_found=False for loader, module_name, is_pkg in pkgutil.iter_modules(modules.__path__): diff --git a/pupy/pupylib/PythonCompleter.py b/pupy/pupylib/PythonCompleter.py index 2d0a6f7c..91109c1e 100644 --- a/pupy/pupylib/PythonCompleter.py +++ b/pupy/pupylib/PythonCompleter.py @@ -1,8 +1,8 @@ import __builtin__ -__all__ = ["PupyCompleter"] +__all__ = ["PythonCompleter"] -class PupyCompleter: +class PythonCompleter: def __init__(self, local_ns=None, global_ns=None): if local_ns is not None: self.local_ns=local_ns @@ -100,7 +100,7 @@ def get_class_members(klass): if __name__=="__main__": import code import readline - readline.set_completer(PupyCompleter().complete) + readline.set_completer(PythonCompleter().complete) readline.parse_and_bind('tab: complete') code.interact()