Add completion for :set command

core/module.py

@@ -11,18 +11,16 @@
 
 """
 
-from core.vectorlist import VectorList
-from core.vectors import ModuleExec
-from core.weexceptions import DevException, ArgparseError
-from core.loggers import log
+import shlex
+
+import utils
 from core import argparsers
 from core import messages
-from mako.template import Template
 from core import modules
-import shlex
-import utils
-import ast
-import os
+from core.loggers import log
+from core.vectorlist import VectorList
+from core.weexceptions import DevException, ArgparseError
+
 
 class Status:
     """Represent the module statuses.
@@ -43,6 +41,7 @@ class Status:
 class Module:
 
     aliases = []
+    arguments = []
 
     def __init__(self, session, name, folder):
         """Module object constructor.
@@ -73,9 +72,53 @@ def __init__(self, session, name, folder):
 
         # Arguments dictionary is initially empty
         self.args = {}
-
         self.init()
 
+        self._register_arguments(self.arguments)
+
+    def complete(self, text) -> []:
+        names = []
+
+        for a in self.arguments:
+            name = a.get('name')
+            if isinstance(name, list):
+                name = name[0]
+            name = name.strip('-')
+            if name.startswith(text):
+                names.append(f'{name} ')
+
+        return names
+
+    def register_arguments(self, arguments):
+        self.arguments = arguments
+
+    def _register_arguments(self, arguments):
+        """Register the module arguments.
+
+        Register arguments to be added to the argparse parser.
+
+        Args:
+            arguments (list of dict): List of dictionaries in the form
+            `[{ 'name' : 'arg1', 'opt' : '', .. }, {'name' : 'arg2', 'opt' : '', .. }]`
+            to be passed to the `ArgumentParser.add_argument()` method.
+        """
+
+        try:
+            for arg_opts in arguments:
+
+                # Handle if the argument registration is done before
+                # The vector registration. This should at least warn
+                if arg_opts.get('choices') == []:
+                    log.warn(messages.module.error_choices_s_s_empty % (self.name,
+                                                                        arg_name))
+
+                self.argparser.add_argument(
+                    arg_opts['name'],
+                    **dict((k, v) for k, v in arg_opts.items() if k != 'name')
+                )
+        except Exception as e:
+            raise DevException(messages.module.error_setting_arguments_s % (e))
+
     def run_cmdline(self, line, cmd = ''):
         """Execute the module from command line.
 
@@ -306,33 +349,6 @@ def register_info(self, info):
         if not self.argparser.description:
             raise DevException(messages.module.error_module_missing_description)
 
-    def register_arguments(self, arguments = []):
-        """Register the module arguments.
-
-        Register arguments to be added to the argparse parser.
-
-        Args:
-            arguments (list of dict): List of dictionaries in the form
-            `[{ 'name' : 'arg1', 'opt' : '', .. }, {'name' : 'arg2', 'opt' : '', .. }]`
-            to be passed to the `ArgumentParser.add_argument()` method.
-        """
-
-        try:
-            for arg_opts in arguments:
-
-                # Handle if the argument registration is done before
-                # The vector registration. This should at least warn
-                if arg_opts.get('choices') == []:
-                    log.warn(messages.module.error_choices_s_s_empty % (self.name,
-                                                                        arg_name))
-
-                self.argparser.add_argument(
-                    arg_opts['name'],
-                    **dict((k, v) for k, v in arg_opts.items() if k != 'name')
-                )
-        except Exception as e:
-            raise DevException(messages.module.error_setting_arguments_s % (e))
-
 
     def register_vectors(self, vectors):
         """Register the module vectors.

core/sessions.py

@@ -1,28 +1,26 @@
-from core import messages
-from core.weexceptions import FatalException
-from mako import template
-from core.config import sessions_path, sessions_ext
-from core.loggers import log, dlog, stream_handler
-from core.module import Status
-import os
-import yaml
+import ast
+import atexit
 import glob
 import logging
-import urllib.parse
-import atexit
-import ast
+import os
 import pprint
+import urllib.parse
 
-print_filters = (
-    'debug',
-    'channel',
-    'proxy'
-)
+import yaml
+from mako import template
 
-set_filters = (
-    'debug',
-    'channel',
-    'proxy'
+from core import messages
+from core.config import sessions_path, sessions_ext
+from core.loggers import log, dlog, stream_handler
+from core.module import Status
+from core.weexceptions import FatalException
+
+arg_blacklist = (
+    'default_shell',
+    'password',
+    'command',
+    'path',
+    'url',
 )
 
 class Session(dict):
@@ -37,22 +35,34 @@ def _session_save_atexit(self):
             )
 
     def print_to_user(self, module_filter = ''):
-
         dlog.info(pprint.pformat(self))
 
-        for mod_name, mod_value in self.items():
+        for arg in self.get_stored_args(module_filter):
+            log.info(messages.sessions.set_s_s % arg)
 
-            if isinstance(mod_value, dict):
-                mod_args = mod_value.get('stored_args')
+    def get_stored_args(self, term = ''):
+        args = []
 
-                # Is a module, print all the storable stored_arguments
-                for argument, arg_value in mod_args.items():
-                    if not module_filter or ("%s.%s" % (mod_name, argument)).startswith(module_filter):
-                        log.info(messages.sessions.set_module_s_s_s % (mod_name, argument, arg_value))
-            else:
-                # If is not a module, just print if matches with print_filters
-                if any(f for f in print_filters if f == mod_name and f.startswith(module_filter)):
-                    log.info(messages.sessions.set_s_s % (mod_name, mod_value))
+        for mod_name, mod_value in self.items():
+            # Is a module, print all the storable stored_arguments
+            if isinstance(mod_value, dict):
+                for argument, arg_value in mod_value.get('stored_args', {}).items():
+                    path = ("%s.%s" % (mod_name, argument))
+                    if not term or path.startswith(term):
+                        args.append((path, arg_value))
+            # If is not a module, just print if not in blacklist
+            elif mod_name not in arg_blacklist and mod_name.startswith(term):
+                args.append((mod_name, mod_value))
+
+        return args
+
+    def complete(self, text):
+        names = []
+        for name, val in self.items():
+            if not isinstance(val, dict) and name not in arg_blacklist and name.startswith(text):
+                names.append(name)
+
+        return names
 
     def get_connection_info(self):
      return template.Template(messages.sessions.connection_info).render(
@@ -119,7 +129,7 @@ def unset(self, module_argument):
                 log.info(messages.sessions.unset_module_s_s % (module_name, arg_name))
         else:
             module_name = module_argument
-            if module_name not in self and module_name not in set_filters:
+            if module_name in arg_blacklist or module_name not in self:
                 log.warning(messages.sessions.error_session_s_not_modified % (module_name))
             else:
                 self[module_name] = None
@@ -155,7 +165,7 @@ def set(self, module_argument, value):
             log.info(messages.sessions.set_module_s_s_s % (module_name, arg_name, value))
         else:
             module_name = module_argument
-            if module_name not in self and module_name not in set_filters:
+            if module_name in arg_blacklist or module_name not in self:
                 log.warning(messages.sessions.error_session_s_not_modified % (module_name))
             else:
                 self[module_name] = value
@@ -264,6 +274,7 @@ def __init__(self, url, password, volatile = False):
                         'debug': False,
                         'channel': None,
                         'default_shell': None,
+                        'proxy': None,
                     }
                 )
 

core/terminal.py

@@ -1,23 +1,22 @@
-from core.weexceptions import FatalException, ChannelException
-from core.loggers import log, dlog
+from mako import template
+
+import utils
+from core import config
 from core import messages
 from core import modules
-from core import config
+from core.loggers import log, dlog
 from core.module import Status
-import utils
-from mako import template
+from core.weexceptions import ChannelException
 
 try:
     import gnureadline as readline
 except ImportError:
     import readline
 
 import cmd
-import glob
-import os
 import shlex
 import atexit
-import sys
+
 
 class CmdModules(cmd.Cmd):
 
@@ -286,6 +285,26 @@ def do_show(self, line, cmd):
 
         self.session.print_to_user(line)
 
+    def completenames_and_args(self, text):
+        names = []
+
+        mod_text, arg_text = text.split('.') if '.' in text else (text, '',)
+
+        for n in self.completenames(mod_text):
+            if n.startswith(mod_text):
+                if n in modules.loaded:
+                    mod = modules.loaded[n]
+                    if mod_text == n:
+                        names.extend([f'{n}.{o}' for o in mod.complete(arg_text)])
+                    elif len(mod.arguments):
+                        names.append(f'{n}.')
+        return names
+
+    def complete_set(self, text, line, begidx, endidx):
+        if line.count(' ') < 2:
+            res = self.session.complete(text) + self.completenames_and_args(text)
+            return res
+
     def do_set(self, line, cmd):
         """Command "set" to set session variables."""
 

modules/audit/filesystem.py

@@ -1,13 +1,13 @@
-from core.vectors import PhpCode, ShellCmd, ModuleExec, Os
-from core.module import Module
 from core.loggers import log
-from core import modules
-import utils
+from core.module import Module
+from core.vectors import ModuleExec
 
-class Filesystem(Module):
 
+class Filesystem(Module):
     """Audit the file system for weak permissions."""
 
+    check_functions = []
+
     def init(self):
 
         self.register_info(