From acd5f58781d2ec21a03f1709088b8b14d819a3b3 Mon Sep 17 00:00:00 2001
From: Lucien A <lu.aubert84@gmail.com>
Date: Fri, 16 Jun 2023 18:01:41 +0200
Subject: [PATCH 1/4] replace cmd.Cmd by prompt_toolkit

---
 core/messages.py     |  24 +--
 core/old_terminal.py | 355 ++++++++++++++++++++++++++++++++++
 core/sessions.py     |  33 ++--
 core/terminal.py     | 440 +++++++++++++++++++++----------------------
 modules/shell/sh.py  |  10 +-
 requirements.txt     |   1 +
 6 files changed, 594 insertions(+), 269 deletions(-)
 create mode 100644 core/old_terminal.py

diff --git a/core/messages.py b/core/messages.py
index 0e420963..1c56e88e 100644
--- a/core/messages.py
+++ b/core/messages.py
@@ -29,18 +29,6 @@ class sessions:
     unset_s = '%s is now unset'
     error_loading_sessions = 'Session loading error'
     error_session_s_not_modified = 'Error setting session variable \'%s\''
-    connection_info = """<%!
-from urllib.parse import urlparse
-%><%
-if not host:
-    urlparsed = urlparse(url)
-    if urlparsed and urlparsed.netloc:
-        hostname = urlparsed.netloc
-    else:
-        hostname = 'undefined host'
-else:
-    hostname = host
-%>${'%s@' % user if user else ''}${hostname}${':%s' % path if path and path != '.' else ''}"""
 
 class channels:
     error_loading_channel_s = 'Error loading channel \'%s\''
@@ -58,16 +46,16 @@ class terminal:
 following command replacements to simulate a unrestricted shell.
 """
     welcome_to_s = """
-[+] weevely ${version}
+<gutter>[+]</gutter> <label>Weevely</label> <underline><value>${version}</value></underline>
 
-[+] Target:\t${conn_info}
-[+] Session:\t${path}
+<gutter>[+]</gutter> <label>Target:</label>\t<value>${'%s@' % user if user else ''}${hostname}${':%s' % conn_path if conn_path and conn_path != '.' else ''}</value>
+<gutter>[+]</gutter> <label>Session:</label>\t<value>${path}</value>
 % if default_shell:
-[+] Shell:\t${ 'System shell' if default_shell == 'shell_sh' else 'PHP interpreter'}
+<gutter>[+]</gutter> <label>Shell:</label>\t<value>${ 'System shell' if default_shell == 'shell_sh' else 'PHP interpreter'}</value>
 % endif
 
-[+] Browse the filesystem or execute commands starts the connection
-[+] to the target. Type :help for more information.
+<gutter>[+]</gutter> Browse the filesystem or execute commands starts the connection
+<gutter>[+]</gutter> to the target. Type :help for more information.
 """
     set_usage = 'Set session variable (run :show to print). Usage:\n:set <variable> \'<value>\''
     unset_usage = 'Unset session variable (run :show to print). Usage:\n:unset <variable>'
diff --git a/core/old_terminal.py b/core/old_terminal.py
new file mode 100644
index 00000000..438a4aee
--- /dev/null
+++ b/core/old_terminal.py
@@ -0,0 +1,355 @@
+from mako import template
+
+import utils
+from core import config
+from core import messages
+from core import modules
+from core.loggers import log, dlog
+from core.module import Status
+from core.weexceptions import ChannelException
+
+try:
+    import gnureadline as readline
+except ImportError:
+    import readline
+
+import cmd
+import shlex
+import atexit
+
+
+class CmdModules(cmd.Cmd):
+
+    identchars = cmd.Cmd.identchars + ':'
+    doc_header = "Modules and commands (type :help <module>):"
+    nohelp = "[!] No help on %s"
+
+    def complete(self, text, state):
+        """Return the next possible completion for 'text'.
+
+        If a command has not been entered, then complete against command list.
+        Otherwise try to call complete_<command> to get list of completions.
+        """
+        if state == 0:
+            origline = readline.get_line_buffer()
+
+            # Offer completion just for commands that starts
+            # with the trigger :
+            if origline and not origline.startswith(':'):
+                return None
+
+            line = origline.lstrip().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)
+                    except AttributeError:
+                        compfunc = self.completedefault
+            else:
+                compfunc = self.completenames
+            self.completion_matches = compfunc(text, line, begidx, endidx)
+        try:
+            if self.completion_matches[state].startswith('alias_'):
+                if self.session.get('default_shell') == 'shell_php':
+                    return self.completion_matches[state][6:]
+                else:
+                    return ''
+            else:
+                return self.completion_matches[state]
+        except IndexError:
+            return None
+
+    def onecmd(self, line):
+        """Interpret the argument as though it had been typed in response
+        to the prompt.
+
+        This may be overridden, but should not normally need to be;
+        see the precmd() and postcmd() methods for useful execution hooks.
+        The return value is a flag indicating whether interpretation of
+        commands by the interpreter should stop.
+
+        """
+        cmd, arg, line = self.parseline(line)
+        if not line:
+            return self.emptyline()
+        if cmd in (None, ''):
+            return self.default(line)
+        self.lastcmd = line
+        if line == 'EOF' :
+            #self.lastcmd = ''
+            raise EOFError()
+        if cmd:
+            # Try running module  command
+            try:
+                func = getattr(self, 'do_' + cmd.lstrip(':'))
+            except AttributeError:
+                # If there is no module command, check if we have a PHP shelli
+                # And in case try running alias command
+                if self.session.get('default_shell') == 'shell_php' or cmd.lstrip(':') == 'cd':
+                    try:
+                        func = getattr(self, 'do_alias_' + cmd.lstrip(':'))
+                    except AttributeError:
+                        pass
+                    else:
+                        return func(arg, cmd)
+            else:
+                return func(arg, cmd)
+
+        return self.default(line)
+
+    def _print_modules(self):
+
+        data = []
+        for module_group, names in modules.loaded_tree.items():
+            for module_name in names:
+                data.append([ ':%s' % module_name, modules.loaded[module_name].info.get('description', '') ])
+
+        if data: log.info(utils.prettify.tablify(data, table_border = False))
+
+    def _print_command_replacements(self):
+
+        data = []
+        for module_name, module in modules.loaded.items():
+            if module.aliases:
+                data.append([ ', '.join(module.aliases), module_name ])
+
+        if data: log.info(utils.prettify.tablify(data, table_border = False))
+
+    def do_help(self, arg, command):
+        """Fixed help."""
+
+        print()
+
+        self._print_modules()
+
+        if self.session['shell_sh']['status'] == Status.RUN: print(); return
+
+        log.info(messages.terminal.help_no_shell)
+        self._print_command_replacements()
+
+        print()
+
+
+class Terminal(CmdModules):
+
+    """Weevely Terminal"""
+
+    def __init__(self, session):
+
+        cmd.Cmd.__init__(self)
+
+        self.session = session
+        self.prompt = 'weevely> '
+
+        # Load all available modules
+        self._load_modules()
+
+        # Load history file
+        self._load_history()
+
+        # Set a nice intro
+        self.intro = template.Template(
+            messages.terminal.welcome_to_s
+        ).render(
+            path = self.session.get('path'),
+            conn_info = session.get_connection_info(),
+            version = messages.version,
+            default_shell = self.session.get('default_shell')
+        )
+
+    def emptyline(self):
+        """Disable repetition of last command."""
+
+        pass
+
+    def precmd(self, line):
+        """Before to execute a line commands. Confirm shell availability and get basic system infos """
+
+        dlog.info('>>>> %s' % line)
+
+        # Skip slack check is not a remote command
+        if not line or any(
+                        line.startswith(cmnd) for cmnd in (
+                            ':set',
+                            ':unset',
+                            ':show',
+                            ':help'
+                        )
+                    ):
+            return line
+
+        # Trigger the shell_sh/shell_php probe if
+        # 1. We never tried to raise shells (shell_sh = IDLE)
+        # 2. The basic intepreter shell_php is not running.
+        if (
+            self.session['shell_sh']['status'] == Status.IDLE or
+            self.session['shell_php']['status'] != Status.RUN
+            ):
+
+            # We're implying that no shell is set, so reset default shell
+            self.session['default_shell'] = None
+
+            # Force shell_php to idle to avoid to be skipped by shell_sh
+            self.session['shell_php']['status'] = Status.IDLE
+
+            # Catch every exception which prevent the shell setup.
+            # We imply that at every channel change (proxy, channel name)
+            # this piece of code will be executed.
+            try:
+                self.session['shell_sh']['status'] = modules.loaded['shell_sh'].setup()
+            except ChannelException as e:
+                log.error(str(e))
+                return ''
+
+        # Set default_shell in any case (could have been changed runtime)
+        for shell in ('shell_sh', 'shell_php'):
+
+            if self.session[shell]['status'] == Status.RUN:
+                self.session['default_shell'] = shell
+                break
+
+        # Kill the execution if no shell were loaded
+        if not self.session.get('default_shell'):
+            log.error(messages.terminal.backdoor_unavailable)
+            return ''
+
+        # TODO: do not print this every loop
+        # Print an introductory string with php shell
+        #if self.session.get('default_shell') == 'shell_php':
+        #    log.info(messages.terminal.welcome_no_shell)
+        #    self._print_command_replacements()
+        #    log.info('\nweevely> %s' % line)
+
+        # Get hostname and whoami if not set
+        if not self.session['system_info']['results'].get('hostname'):
+            modules.loaded['system_info'].run_argv([ "-info", "hostname"])
+
+        if not self.session['system_info']['results'].get('whoami'):
+            modules.loaded['system_info'].run_argv(["-info", "whoami"])
+
+        # Get current working directory if not set
+        # Should be OK to repeat this every time if not set.
+        if not self.session['file_cd']['results'].get('cwd'):
+            self.do_file_cd(".")
+
+        return line
+
+    def postcmd(self, stop, line):
+
+        default_shell = self.session.get('default_shell')
+
+        if not default_shell:
+            self.prompt = 'weevely> '
+        else:
+            if default_shell == 'shell_sh':
+                prompt = '$'
+            elif default_shell == 'shell_php':
+                prompt = 'PHP>'
+            else:
+                prompt = '?'
+
+            # Build next prompt, last command could have changed the cwd
+            self.prompt = '%s %s ' % (self.session.get_connection_info(), prompt)
+
+
+    def default(self, line):
+        """Default command line send."""
+
+        if not line: return
+
+        default_shell = self.session.get('default_shell')
+
+        if not default_shell: return
+
+        result = modules.loaded[default_shell].run_argv([line])
+
+        if not result: return
+
+        # Clean trailing newline if existent to prettify output
+        result = result[:-1] if (
+                isinstance(result, str) and
+                result.endswith('\n')
+            ) else result
+
+        log.info(result)
+
+    def do_show(self, line, cmd):
+        """Command "show" which prints session variables"""
+
+        self.session.print_to_user(line)
+
+    def do_set(self, line, cmd):
+        """Command "set" to set session variables."""
+
+        try:
+            args = shlex.split(line)
+        except Exception as e:
+            import traceback; log.debug(traceback.format_exc())
+            log.warning(messages.generic.error_parsing_command_s % str(e))
+
+        # Set the setting
+        else:
+            if len(args) < 2:
+                log.warning(messages.terminal.set_usage)
+            elif len(args) >= 2:
+                args[1] = ' '.join(args[1:])
+                self.session.set(args[0], args[1])
+
+    def do_unset(self, line, cmd):
+        """Command "unset" to unset session variables."""
+
+        # Print all settings that startswith args[0]
+        if not line:
+            log.warning(messages.terminal.unset_usage)
+
+        # Set the setting
+        else:
+            self.session.unset(line)
+
+    def _load_modules(self):
+        """Load all modules assigning corresponding do_* functions."""
+
+        for module_name, module_class in modules.loaded.items():
+
+            # Set module.do_terminal_module() function as terminal
+            # self.do_modulegroup_modulename()
+            setattr(
+                Terminal, 'do_%s' %
+                (module_name), module_class.run_cmdline)
+
+            # Set module.do_alias() function as terminal
+            # self.do_alias() for every defined `Module.aliases`.
+            for alias in module_class.aliases:
+                setattr(
+                    Terminal, 'do_alias_%s' %
+                    (alias), module_class.run_alias)
+                setattr(
+                    Terminal, 'help_%s' %
+                    (alias), module_class.help)
+
+            # Set module.help() function as terminal
+            # self.help_modulegroup_modulename()
+            setattr(
+                Terminal, 'help_%s' %
+                (module_name), module_class.help)
+
+    def _load_history(self):
+        """Load history file and register dump on exit."""
+
+        # Create a file without truncating it in case it exists.
+        open(config.history_path, 'a').close()
+
+        readline.set_history_length(100)
+        try:
+            readline.read_history_file(config.history_path)
+        except IOError:
+            pass
+        atexit.register(readline.write_history_file,
+            config.history_path)
diff --git a/core/sessions.py b/core/sessions.py
index b9c90a36..7fb854c3 100644
--- a/core/sessions.py
+++ b/core/sessions.py
@@ -1,17 +1,18 @@
+import ast
+import atexit
+import glob
+import logging
+import os
+import pprint
+import urllib.parse
+
+import yaml
+
 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 glob
-import logging
-import urllib.parse
-import atexit
-import ast
-import pprint
+from core.weexceptions import FatalException
 
 print_filters = (
     'debug',
@@ -55,12 +56,12 @@ def print_to_user(self, module_filter = ''):
                     log.info(messages.sessions.set_s_s % (mod_name, mod_value))
 
     def get_connection_info(self):
-     return template.Template(messages.sessions.connection_info).render(
-         url = self['url'],
-         user = self['system_info']['results'].get('whoami', ''),
-         host = self['system_info']['results'].get('hostname', ''),
-         path = self['file_cd']['results'].get('cwd', '.')
-     )
+        return {
+            'url': self['url'],
+            'user': self['system_info']['results'].get('whoami', ''),
+            'host': self['system_info']['results'].get('hostname', ''),
+            'path': self['file_cd']['results'].get('cwd', '.')
+        }
 
     def load_session(self, data):
         """
diff --git a/core/terminal.py b/core/terminal.py
index 8860fea7..53b40173 100644
--- a/core/terminal.py
+++ b/core/terminal.py
@@ -1,198 +1,168 @@
-from core.weexceptions import FatalException, ChannelException
-from core.loggers import log, dlog
-from core import messages
-from core import modules
-from core import config
-from core.module import Status
-import utils
-from mako import template
-
-try:
-    import gnureadline as readline
-except ImportError:
-    import readline
+import string
+from urllib.parse import urlparse
 
-import cmd
-import glob
-import os
-import shlex
-import atexit
-import sys
+from mako import template
+from prompt_toolkit import PromptSession, HTML
+from prompt_toolkit import print_formatted_text
+from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.history import FileHistory
+from prompt_toolkit.key_binding import KeyBindings
+from prompt_toolkit.key_binding.bindings import named_commands
+from prompt_toolkit.output import ColorDepth
+from prompt_toolkit.styles import Style
+
+from core import messages, modules, config
+from core.loggers import log
+from core.module import Status
+from core.weexceptions import ChannelException
+
+IDENTCHARS = string.ascii_letters + string.digits + '_'
+
+
+class CustomCompleter(Completer):
+    terminal = dict()
+
+    def __init__(self, terminal):
+        self.terminal = terminal
+
+    def get_completions(self, document, complete_event):
+        for attr in dir(self.terminal):
+            module = ''
+            alias = False
+            if attr.startswith('do_alias_'):
+                module = ':'+attr[9:]
+                alias = True
+            elif attr.startswith('do_'):
+                module = ':'+attr[3:]
+
+            classname = ('alias' if alias else 'module')
+            if module and module.startswith(document.text):
+                yield Completion(module,
+                                 start_position=-document.cursor_position,
+                                 style=f'class:completion-menu.completion.{classname}',
+                                 selected_style=f'class:completion-menu.completion.current.{classname}',
+                                 )
+
+
+class Terminal:
+    session = dict()
+    kb = KeyBindings()
+    style = Style.from_dict({
+        'bottom-toolbar': 'noreverse',
+        'bottom-toolbar.text': 'bg:#ff0000',
+        'rprompt': 'bg:#ff0066 #000000',
+        'username': '#999999 bold',
+        'at': '#00aa00',
+        'colon': '#ffffff',
+        'pound': '#00aa00',
+        'host': '#ff00ff',
+        'path': 'ansicyan',
+        'mode': '#ff0000',
+        'gutter': '#999999',
+        'label': '#00aa00 bold',
+        'value': '#33ff66',
+
+        'completion-menu': 'bg:#96d0f7 #000000',
+        'completion-menu.completion.current': 'bg:#0077ff #000000',
+        'completion-menu.completion.alias': 'bg:#1a9cf2',
+    })
 
-class CmdModules(cmd.Cmd):
+    def __init__(self, session):
+        self.session = session
+        self.prompt_session = PromptSession(history=FileHistory(config.history_path))
+        self.completer = CustomCompleter(self)
 
-    identchars = cmd.Cmd.identchars + ':'
-    doc_header = "Modules and commands (type :help <module>):"
-    nohelp = "[!] No help on %s"
+        self._load_modules()
+        self._print_intro()
 
-    def complete(self, text, state):
-        """Return the next possible completion for 'text'.
+    def cmdloop(self):
+        while True:
+            try:
+                line = self.prompt_session.prompt(self.get_prompt_message,
+                                                  color_depth=ColorDepth.TRUE_COLOR,
+                                                  complete_while_typing=True,
+                                                  reserve_space_for_menu=10,
+                                                  completer=self.completer,
+                                                  key_bindings=self.kb,
+                                                  style=self.style,
+                                                  )
+                line = self.precmd(line)
+                self.onecmd(line)
+            except KeyboardInterrupt:
+                # Quit when pressing Ctrl-C while prompt is empty
+                if len(self.prompt_session.default_buffer.text) == 0:
+                    raise EOFError
 
-        If a command has not been entered, then complete against command list.
-        Otherwise try to call complete_<command> to get list of completions.
-        """
-        if state == 0:
-            origline = readline.get_line_buffer()
-
-            # Offer completion just for commands that starts
-            # with the trigger :
-            if origline and not origline.startswith(':'):
-                return None
-
-            line = origline.lstrip().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)
-                    except AttributeError:
-                        compfunc = self.completedefault
-            else:
-                compfunc = self.completenames
-            self.completion_matches = compfunc(text, line, begidx, endidx)
-        try:
-            if self.completion_matches[state].startswith('alias_'):
-                if self.session.get('default_shell') == 'shell_php':
-                    return self.completion_matches[state][6:]
-                else:
-                    return ''
-            else:
-                return self.completion_matches[state]
-        except IndexError:
-            return None
+    def precmd(self, line):
+        self.init_default_shell()
+        return line
 
     def onecmd(self, line):
-        """Interpret the argument as though it had been typed in response
-        to the prompt.
-
-        This may be overridden, but should not normally need to be;
-        see the precmd() and postcmd() methods for useful execution hooks.
-        The return value is a flag indicating whether interpretation of
-        commands by the interpreter should stop.
-
-        """
-        cmd, arg, line = self.parseline(line)
         if not line:
-            return self.emptyline()
-        if cmd in (None, ''):
-            return self.default(line)
-        self.lastcmd = line
-        if line == 'EOF' :
-            #self.lastcmd = ''
-            raise EOFError()
-        if cmd:
-            # Try running module  command
-            try:
-                func = getattr(self, 'do_' + cmd.lstrip(':'))
-            except AttributeError:
-                # If there is no module command, check if we have a PHP shelli
-                # And in case try running alias command
-                if self.session.get('default_shell') == 'shell_php' or cmd.lstrip(':') == 'cd':
-                    try:
-                        func = getattr(self, 'do_alias_' + cmd.lstrip(':'))
-                    except AttributeError:
-                        pass
-                    else:
-                        return func(arg, cmd)
-            else:
-                return func(arg, cmd)
-
-        return self.default(line)
+            return
 
-    def _print_modules(self):
+        cmd, args, line = self.parseline(line)
 
-        data = []
-        for module_group, names in modules.loaded_tree.items():
-            for module_name in names:
-                data.append([ ':%s' % module_name, modules.loaded[module_name].info.get('description', '') ])
-
-        if data: log.info(utils.prettify.tablify(data, table_border = False))
-
-    def _print_command_replacements(self):
-
-        data = []
-        for module_name, module in modules.loaded.items():
-            if module.aliases:
-                data.append([ ', '.join(module.aliases), module_name ])
-
-        if data: log.info(utils.prettify.tablify(data, table_border = False))
-
-    def do_help(self, arg, command):
-        """Fixed help."""
-
-        print()
-
-        self._print_modules()
-
-        if self.session['shell_sh']['status'] == Status.RUN: print(); return
-
-        log.info(messages.terminal.help_no_shell)
-        self._print_command_replacements()
-
-        print()
-
-
-class Terminal(CmdModules):
-
-    """Weevely Terminal"""
-
-    def __init__(self, session):
+        try:
+            func = getattr(self, 'do_' + cmd.lstrip(':'))
+            return func(args, cmd)
+        except AttributeError:
+            return self.default(line)
 
-        cmd.Cmd.__init__(self)
+    def default(self, line):
+        if not line:
+            return
 
-        self.session = session
-        self.prompt = 'weevely> '
+        default_shell = self.session.get('default_shell')
 
-        # Load all available modules
-        self._load_modules()
+        if not default_shell:
+            return
 
-        # Load history file
-        self._load_history()
+        result = modules.loaded[default_shell].run_argv([line])
 
-        # Set a nice intro
-        self.intro = template.Template(
-            messages.terminal.welcome_to_s
-        ).render(
-            path = self.session.get('path'),
-            conn_info = session.get_connection_info(),
-            version = messages.version,
-            default_shell = self.session.get('default_shell')
-        )
+        if not result:
+            return
 
-    def emptyline(self):
-        """Disable repetition of last command."""
+        # Clean trailing newline if existent to prettify output
+        result = result[:-1] if (
+                isinstance(result, str) and
+                result.endswith('\n')
+        ) else result
 
-        pass
+        print(result)
 
-    def precmd(self, line):
-        """Before to execute a line commands. Confirm shell availability and get basic system infos """
+    def parseline(self, line):
+        """Parse the line into a command name and a string containing
+        the arguments.  Returns a tuple containing (command, args, line).
+        'command' and 'args' may be None if the line couldn't be parsed.
+        """
+        line = line.strip()
 
-        dlog.info('>>>> %s' % line)
+        if line and line[0] == ':':
+            line = line[1:]
 
-        # Skip slack check is not a remote command
-        if not line or any(
-                        line.startswith(cmnd) for cmnd in (
-                            ':set',
-                            ':unset',
-                            ':show',
-                            ':help'
-                        )
-                    ):
-            return line
+        if not line:
+            return None, None, line
+        elif line[0] == '?':
+            line = 'help ' + line[1:]
+        elif line[0] == '!':
+            if hasattr(self, 'do_shell'):
+                line = 'shell ' + line[1:]
+            else:
+                return None, None, line
+        i, n = 0, len(line)
+        while i < n and line[i] in IDENTCHARS: i = i + 1
+        cmd, arg = line[:i], line[i:].strip()
+        return cmd, arg, line
 
+    def init_default_shell(self):
         # Trigger the shell_sh/shell_php probe if
         # 1. We never tried to raise shells (shell_sh = IDLE)
         # 2. The basic intepreter shell_php is not running.
         if (
-            self.session['shell_sh']['status'] == Status.IDLE or
-            self.session['shell_php']['status'] != Status.RUN
-            ):
+                self.session['shell_sh']['status'] == Status.IDLE or
+                self.session['shell_php']['status'] != Status.RUN
+        ):
 
             # We're implying that no shell is set, so reset default shell
             self.session['default_shell'] = None
@@ -223,14 +193,14 @@ def precmd(self, line):
 
         # TODO: do not print this every loop
         # Print an introductory string with php shell
-        #if self.session.get('default_shell') == 'shell_php':
+        # if self.session.get('default_shell') == 'shell_php':
         #    log.info(messages.terminal.welcome_no_shell)
         #    self._print_command_replacements()
         #    log.info('\nweevely> %s' % line)
 
         # Get hostname and whoami if not set
         if not self.session['system_info']['results'].get('hostname'):
-            modules.loaded['system_info'].run_argv([ "-info", "hostname"])
+            modules.loaded['system_info'].run_argv(["-info", "hostname"])
 
         if not self.session['system_info']['results'].get('whoami'):
             modules.loaded['system_info'].run_argv(["-info", "whoami"])
@@ -240,46 +210,38 @@ def precmd(self, line):
         if not self.session['file_cd']['results'].get('cwd'):
             self.do_file_cd(".")
 
-        return line
-
-    def postcmd(self, stop, line):
-
-        default_shell = self.session.get('default_shell')
-
-        if not default_shell:
-            self.prompt = 'weevely> '
-        else:
-            if default_shell == 'shell_sh':
-                prompt = '$'
-            elif default_shell == 'shell_php':
-                prompt = 'PHP>'
-            else:
-                prompt = '?'
-
-            # Build next prompt, last command could have changed the cwd
-            self.prompt = '%s %s ' % (self.session.get_connection_info(), prompt)
-
-
-    def default(self, line):
-        """Default command line send."""
-
-        if not line: return
-
-        default_shell = self.session.get('default_shell')
-
-        if not default_shell: return
-
-        result = modules.loaded[default_shell].run_argv([line])
-
-        if not result: return
-
-        # Clean trailing newline if existent to prettify output
-        result = result[:-1] if (
-                isinstance(result, str) and
-                result.endswith('\n')
-            ) else result
-
-        log.info(result)
+    def get_prompt_message(self):
+        shell = self.session.get('default_shell')
+        pound = '?'
+        if not shell:
+            pound = 'weevely>'
+        elif shell == 'shell_sh':
+            pound = '$'
+        elif shell == 'shell_php':
+            pound = 'PHP>'
+
+        info = self.session.get_connection_info()
+
+        if not shell or len(info.get('user')) == 0:
+            return [
+                ('class:pound', pound),
+                ('class:space', ' '),
+            ]
+
+        msg = []
+
+        msg.extend([
+            ('class:username', info.get('user')),
+            ('class:at', '@'),
+            ('class:host', info.get('host')),
+            ('class:colon', ':'),
+            ('class:path', info.get('path')),
+            ('class:space', ' '),
+            ('class:pound', pound),
+            ('class:space', ' '),
+        ])
+
+        return msg
 
     def do_show(self, line, cmd):
         """Command "show" which prints session variables"""
@@ -292,7 +254,8 @@ def do_set(self, line, cmd):
         try:
             args = shlex.split(line)
         except Exception as e:
-            import traceback; log.debug(traceback.format_exc())
+            import traceback;
+            log.debug(traceback.format_exc())
             log.warning(messages.generic.error_parsing_command_s % str(e))
 
         # Set the setting
@@ -321,36 +284,53 @@ def _load_modules(self):
 
             # Set module.do_terminal_module() function as terminal
             # self.do_modulegroup_modulename()
-            setattr(
-                Terminal, 'do_%s' %
-                (module_name), module_class.run_cmdline)
+            setattr(Terminal, 'do_%s' % (module_name), module_class.run_cmdline)
 
             # Set module.do_alias() function as terminal
             # self.do_alias() for every defined `Module.aliases`.
             for alias in module_class.aliases:
                 setattr(
                     Terminal, 'do_alias_%s' %
-                    (alias), module_class.run_alias)
+                              (alias), module_class.run_alias)
                 setattr(
                     Terminal, 'help_%s' %
-                    (alias), module_class.help)
+                              (alias), module_class.help)
 
             # Set module.help() function as terminal
             # self.help_modulegroup_modulename()
             setattr(
                 Terminal, 'help_%s' %
-                (module_name), module_class.help)
+                          (module_name), module_class.help)
 
-    def _load_history(self):
-        """Load history file and register dump on exit."""
+    def _print_intro(self):
+        info = self.session.get_connection_info()
 
-        # Create a file without truncating it in case it exists.
-        open(config.history_path, 'a').close()
+        hostname = info.get('host')
+        if not hostname:
+            urlparsed = urlparse(info.get('url'))
+            if urlparsed and urlparsed.netloc:
+                hostname = urlparsed.netloc
+            else:
+                hostname = 'undefined host'
 
-        readline.set_history_length(100)
-        try:
-            readline.read_history_file(config.history_path)
-        except IOError:
-            pass
-        atexit.register(readline.write_history_file,
-            config.history_path)
+        print_formatted_text(HTML(template.Template(
+            messages.terminal.welcome_to_s
+        ).render(
+            path=self.session.get('path'),
+            version=messages.version,
+            default_shell=self.session.get('default_shell'),
+            url=info.get('url'),
+            user=info.get('user'),
+            hostname=hostname,
+            conn_path=info.get('path'),
+        )), style=self.style)
+
+    @staticmethod
+    @kb.add('enter')
+    def _enter_key(event) -> None:
+        buff = event.current_buffer
+        if buff.complete_state:
+            named_commands.complete(event)
+            buff.insert_text(' ')
+        else:
+            named_commands.accept_line(event)
diff --git a/modules/shell/sh.py b/modules/shell/sh.py
index f1200d44..8c938842 100644
--- a/modules/shell/sh.py
+++ b/modules/shell/sh.py
@@ -1,10 +1,9 @@
-from core.vectors import PhpCode
+import random
+
 from core.module import Module, Status
-from core.loggers import log
 from core.vectors import Os
-from core import messages
-from core import modules
-import random
+from core.vectors import PhpCode
+
 
 class Sh(Module):
 
@@ -46,6 +45,7 @@ def init(self):
                     proc_close($h);
                 }""", "proc_open"),
             PhpCode("""@python_eval('import os; os.system('${command}${stderr_redirection}');');""", "python_eval"),
+            PhpCode("""@system('python -c "import os;os.system('${command}${stderr_redirection}');"');""", "safe_mode_py"),
             PhpCode("""
                 if(class_exists('Perl')){
                     $perl=new Perl();
diff --git a/requirements.txt b/requirements.txt
index db191a80..f66d6d26 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,3 +4,4 @@ PyYAML
 python-dateutil
 PySocks
 pyOpenSSL
+prompt_toolkit

From 617be9e7b28f67eb2722eeda774db44ab2b9a614 Mon Sep 17 00:00:00 2001
From: Lucien A <lu.aubert84@gmail.com>
Date: Fri, 16 Jun 2023 19:45:00 +0200
Subject: [PATCH 2/4] add help formatting

---
 core/loggers.py         | 25 ++++++++++--
 core/messages.py        | 10 ++---
 core/style.py           | 25 ++++++++++++
 core/terminal.py        | 87 +++++++++++++++++++++++------------------
 modules/backdoor/tcp.py | 14 ++++---
 modules/shell/ssh.py    |  2 +-
 6 files changed, 110 insertions(+), 53 deletions(-)
 create mode 100644 core/style.py

diff --git a/core/loggers.py b/core/loggers.py
index cca0621e..edb391a4 100644
--- a/core/loggers.py
+++ b/core/loggers.py
@@ -1,14 +1,18 @@
-import logging.handlers
 import logging
-import sys
-import core.config
+import logging.handlers
 import os
+import xml
+
+from prompt_toolkit import print_formatted_text
+from prompt_toolkit.formatted_text import HTML
+
+import core.config
+from core import style
 
 log = None
 logfile = None
 
 class WeevelyFormatter(logging.Formatter):
-
     FORMATS = {
         # logging.DEBUG :"[D][%(module)s.%(funcName)s:%(lineno)d] %(message)s",
         logging.DEBUG: "[D][%(module)s] %(message)s",
@@ -23,6 +27,18 @@ def format(self, record):
         return logging.Formatter.format(self, record)
 
 
+class StderrHandler(logging.Handler):
+    def __init__(self):
+        super().__init__()
+        self.formatter = WeevelyFormatter()
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            print_formatted_text(HTML(record.msg), style=style.default_style)
+        except xml.parsers.expat.ExpatError:
+            print(record.msg)
+
+
 if not os.path.isdir(core.config.base_path):
     os.makedirs(core.config.base_path)
 
@@ -40,6 +56,7 @@ def format(self, record):
 
 """Initialize the normal handler"""
 stream_handler = logging.StreamHandler()
+stream_handler = StderrHandler()
 stream_handler.setFormatter(WeevelyFormatter())
 
 """Initialize the standard logger"""
diff --git a/core/messages.py b/core/messages.py
index 1c56e88e..cdf459d4 100644
--- a/core/messages.py
+++ b/core/messages.py
@@ -42,11 +42,11 @@ class terminal:
 command replacements to simulate an unrestricted shell.
 """
     help_no_shell = """
-The system shell interpreter is not available in this session, use the
-following command replacements to simulate a unrestricted shell.
+<warning>The system shell interpreter is not available in this session, use the
+following command replacements to simulate a unrestricted shell.</warning>
 """
     welcome_to_s = """
-<gutter>[+]</gutter> <label>Weevely</label> <underline><value>${version}</value></underline>
+    <label>Weevely</label> <underline><value>${version}</value></underline>
 
 <gutter>[+]</gutter> <label>Target:</label>\t<value>${'%s@' % user if user else ''}${hostname}${':%s' % conn_path if conn_path and conn_path != '.' else ''}</value>
 <gutter>[+]</gutter> <label>Session:</label>\t<value>${path}</value>
@@ -54,8 +54,8 @@ class terminal:
 <gutter>[+]</gutter> <label>Shell:</label>\t<value>${ 'System shell' if default_shell == 'shell_sh' else 'PHP interpreter'}</value>
 % endif
 
-<gutter>[+]</gutter> Browse the filesystem or execute commands starts the connection
-<gutter>[+]</gutter> to the target. Type :help for more information.
+  Browse the filesystem or execute commands to initiate the connection to the target.
+  Type :help for more information.
 """
     set_usage = 'Set session variable (run :show to print). Usage:\n:set <variable> \'<value>\''
     unset_usage = 'Unset session variable (run :show to print). Usage:\n:unset <variable>'
diff --git a/core/style.py b/core/style.py
new file mode 100644
index 00000000..1bdcc3f3
--- /dev/null
+++ b/core/style.py
@@ -0,0 +1,25 @@
+from prompt_toolkit.styles import Style
+
+default_style = Style.from_dict({
+    'bottom-toolbar': 'noreverse',
+    'bottom-toolbar.text': 'bg:#ff0000',
+    'rprompt': 'bg:#ff0066 #000000',
+    'username': '#a6e22e bold',
+    'username.root': '#ff4689 underline',
+    'at': '#ae81ff',
+    'colon': '#ae81ff',
+    'pound': '#959077',
+    'host': '#66d9ef',
+    'path': 'ansicyan',
+    'mode': '#ff0000',
+    'gutter': '#999999',
+    'label': '#00aa00 bold',
+    'value': '#33ff66',
+
+    'warning': '#ff4689',
+    'danger': '#ed007e',
+
+    'completion-menu': 'bg:#96d0f7 #000000',
+    'completion-menu.completion.current': 'bg:#0077ff #000000',
+    'completion-menu.completion.alias': 'bg:#1a9cf2',
+})
\ No newline at end of file
diff --git a/core/terminal.py b/core/terminal.py
index 53b40173..add7d904 100644
--- a/core/terminal.py
+++ b/core/terminal.py
@@ -1,17 +1,17 @@
+import shlex
 import string
 from urllib.parse import urlparse
 
 from mako import template
-from prompt_toolkit import PromptSession, HTML
-from prompt_toolkit import print_formatted_text
+from prompt_toolkit import PromptSession
 from prompt_toolkit.completion import Completer, Completion
 from prompt_toolkit.history import FileHistory
 from prompt_toolkit.key_binding import KeyBindings
 from prompt_toolkit.key_binding.bindings import named_commands
 from prompt_toolkit.output import ColorDepth
-from prompt_toolkit.styles import Style
 
-from core import messages, modules, config
+import utils
+from core import messages, modules, config, style
 from core.loggers import log
 from core.module import Status
 from core.weexceptions import ChannelException
@@ -30,10 +30,10 @@ def get_completions(self, document, complete_event):
             module = ''
             alias = False
             if attr.startswith('do_alias_'):
-                module = ':'+attr[9:]
+                module = ':' + attr[9:]
                 alias = True
             elif attr.startswith('do_'):
-                module = ':'+attr[3:]
+                module = ':' + attr[3:]
 
             classname = ('alias' if alias else 'module')
             if module and module.startswith(document.text):
@@ -47,25 +47,6 @@ def get_completions(self, document, complete_event):
 class Terminal:
     session = dict()
     kb = KeyBindings()
-    style = Style.from_dict({
-        'bottom-toolbar': 'noreverse',
-        'bottom-toolbar.text': 'bg:#ff0000',
-        'rprompt': 'bg:#ff0066 #000000',
-        'username': '#999999 bold',
-        'at': '#00aa00',
-        'colon': '#ffffff',
-        'pound': '#00aa00',
-        'host': '#ff00ff',
-        'path': 'ansicyan',
-        'mode': '#ff0000',
-        'gutter': '#999999',
-        'label': '#00aa00 bold',
-        'value': '#33ff66',
-
-        'completion-menu': 'bg:#96d0f7 #000000',
-        'completion-menu.completion.current': 'bg:#0077ff #000000',
-        'completion-menu.completion.alias': 'bg:#1a9cf2',
-    })
 
     def __init__(self, session):
         self.session = session
@@ -84,7 +65,7 @@ def cmdloop(self):
                                                   reserve_space_for_menu=10,
                                                   completer=self.completer,
                                                   key_bindings=self.kb,
-                                                  style=self.style,
+                                                  style=style.default_style,
                                                   )
                 line = self.precmd(line)
                 self.onecmd(line)
@@ -94,6 +75,9 @@ def cmdloop(self):
                     raise EOFError
 
     def precmd(self, line):
+        if line == 'exit':
+            raise EOFError
+
         self.init_default_shell()
         return line
 
@@ -129,7 +113,7 @@ def default(self, line):
                 result.endswith('\n')
         ) else result
 
-        print(result)
+        log.info(result)
 
     def parseline(self, line):
         """Parse the line into a command name and a string containing
@@ -213,25 +197,30 @@ def init_default_shell(self):
     def get_prompt_message(self):
         shell = self.session.get('default_shell')
         pound = '?'
+        host = ''
         if not shell:
-            pound = 'weevely>'
+            host = 'weevely'
+            pound = '>'
         elif shell == 'shell_sh':
             pound = '$'
         elif shell == 'shell_php':
             pound = 'PHP>'
 
         info = self.session.get_connection_info()
+        user = info.get('user')
 
-        if not shell or len(info.get('user')) == 0:
+        if not shell or len(user) == 0:
             return [
+                ('class:host', host),
                 ('class:pound', pound),
                 ('class:space', ' '),
             ]
 
         msg = []
+        userclass = 'class:username' + ('.root' if user == 'root' else '')
 
         msg.extend([
-            ('class:username', info.get('user')),
+            (userclass, user),
             ('class:at', '@'),
             ('class:host', info.get('host')),
             ('class:colon', ':'),
@@ -243,18 +232,23 @@ def get_prompt_message(self):
 
         return msg
 
+    def do_help(self, line, cmd):
+        self._print_modules(line)
+
+        if self.session['shell_sh']['status'] != Status.RUN:
+            log.info(messages.terminal.help_no_shell)
+            self._print_command_replacements()
+
     def do_show(self, line, cmd):
         """Command "show" which prints session variables"""
-
         self.session.print_to_user(line)
 
     def do_set(self, line, cmd):
         """Command "set" to set session variables."""
-
         try:
             args = shlex.split(line)
         except Exception as e:
-            import traceback;
+            import traceback
             log.debug(traceback.format_exc())
             log.warning(messages.generic.error_parsing_command_s % str(e))
 
@@ -268,7 +262,6 @@ def do_set(self, line, cmd):
 
     def do_unset(self, line, cmd):
         """Command "unset" to unset session variables."""
-
         # Print all settings that startswith args[0]
         if not line:
             log.warning(messages.terminal.unset_usage)
@@ -279,7 +272,6 @@ def do_unset(self, line, cmd):
 
     def _load_modules(self):
         """Load all modules assigning corresponding do_* functions."""
-
         for module_name, module_class in modules.loaded.items():
 
             # Set module.do_terminal_module() function as terminal
@@ -313,7 +305,7 @@ def _print_intro(self):
             else:
                 hostname = 'undefined host'
 
-        print_formatted_text(HTML(template.Template(
+        log.info(template.Template(
             messages.terminal.welcome_to_s
         ).render(
             path=self.session.get('path'),
@@ -323,7 +315,28 @@ def _print_intro(self):
             user=info.get('user'),
             hostname=hostname,
             conn_path=info.get('path'),
-        )), style=self.style)
+        ))
+
+    def _print_modules(self, term):
+        for module_group, names in modules.loaded_tree.items():
+            elements = []
+            for module_name in names:
+                if term in module_name:
+                    elements.append([':%s' % module_name, modules.loaded[module_name].info.get('description', '')])
+
+            if not term or len(elements):
+                log.info('<label>%s</label>' % module_group.capitalize())
+                log.info(utils.prettify.tablify(elements, table_border=False))
+                print()
+
+    def _print_command_replacements(self):
+        data = []
+        for module_name, module in modules.loaded.items():
+            if module.aliases:
+                data.append([', '.join(module.aliases), module_name])
+
+        if data:
+            log.info(utils.prettify.tablify(data, table_border=False))
 
     @staticmethod
     @kb.add('enter')
diff --git a/modules/backdoor/tcp.py b/modules/backdoor/tcp.py
index fbe36939..1a35f901 100644
--- a/modules/backdoor/tcp.py
+++ b/modules/backdoor/tcp.py
@@ -1,10 +1,12 @@
-from core.vectors import PythonCode, ShellCmd, Os
-from core.module import Module
-from core.loggers import log
-from core import messages
-import urllib.parse
 import telnetlib
 import time
+import urllib.parse
+
+from core import messages
+from core.loggers import log
+from core.module import Module
+from core.vectors import PythonCode, ShellCmd, Os
+
 
 class Tcp(Module):
     """Spawn a shell on a TCP port."""
@@ -72,7 +74,7 @@ def init(self):
           { 'name' : '-vector', 'choices' : self.vectors.get_names() }
         ])
 
-    def run(self):
+    def run(self, **kwargs):
 
         # Run all the vectors
         for vector in self.vectors:
diff --git a/modules/shell/ssh.py b/modules/shell/ssh.py
index ccd6ba25..780484f8 100644
--- a/modules/shell/ssh.py
+++ b/modules/shell/ssh.py
@@ -12,7 +12,7 @@ class Ssh(Module):
   py.pexpect      |  No    | Rarely installed
   py.subprocess   |  Yes   | Stores clear passwd on disk temporarily
 
-/!\\ When using py.subprocess vector the password is stored in a file at <ASKPASS> /!\\
+/!\\ When using py.subprocess vector the password is stored in a file at ASKPASS /!\\
     As a result the file gets truncated, chmod and then removed. Becareful !
     This file has to be executable, you have to insure it spawns on a partition
     that allows execution (ie. not mounted with `noexec`).

From b2ff3759d3b52742cb2a49dc0fc941acbad51600 Mon Sep 17 00:00:00 2001
From: Lucien A <lu.aubert84@gmail.com>
Date: Fri, 16 Jun 2023 21:25:43 +0200
Subject: [PATCH 3/4] oops

---
 core/__init__.py       |  0
 core/terminal.py       |  2 +-
 tests/test_terminal.py | 34 +++++++++++++++++-----------------
 3 files changed, 18 insertions(+), 18 deletions(-)
 delete mode 100644 core/__init__.py

diff --git a/core/__init__.py b/core/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/core/terminal.py b/core/terminal.py
index add7d904..ecf99ff1 100644
--- a/core/terminal.py
+++ b/core/terminal.py
@@ -54,9 +54,9 @@ def __init__(self, session):
         self.completer = CustomCompleter(self)
 
         self._load_modules()
-        self._print_intro()
 
     def cmdloop(self):
+        self._print_intro()
         while True:
             try:
                 line = self.prompt_session.prompt(self.get_prompt_message,
diff --git a/tests/test_terminal.py b/tests/test_terminal.py
index 93ec072f..19ce03a7 100644
--- a/tests/test_terminal.py
+++ b/tests/test_terminal.py
@@ -1,11 +1,14 @@
-from tests.base_test import BaseTest
-from core.weexceptions import FatalException
+import subprocess
+
 from testfixtures import log_capture
-from core.terminal import Terminal
-from core.sessions import SessionURL, SessionFile
-from core import modules
+
 from core import messages
-import subprocess
+from core import modules
+from core.sessions import SessionURL, SessionFile
+from core.terminal import Terminal
+from core.weexceptions import FatalException
+from tests.base_test import BaseTest
+
 
 def setUpModule():
     subprocess.check_output("""
@@ -34,7 +37,6 @@ def setUp(self, log_captured):
     def _assert_exec(self, line, expected, log_captured):
         line = self.terminal.precmd(line)
         stop = self.terminal.onecmd(line)
-        stop = self.terminal.postcmd(stop, line)
 
         self.assertEqual(log_captured.records[-1].msg, expected)
 
@@ -76,7 +78,6 @@ def test_run_wrong_pass(self, log_captured):
         line = 'echo 1'
         line = terminal.precmd(line)
         stop = terminal.onecmd(line)
-        stop = terminal.postcmd(stop, line)
 
         # Test the behaviour when starting terminal on wrong remote pass
         self.assertTrue(
@@ -95,7 +96,6 @@ def test_run_wrong_url(self, log_captured):
         line = 'echo 1'
         line = terminal.precmd(line)
         stop = terminal.onecmd(line)
-        stop = terminal.postcmd(stop, line)
 
         # Test the behaviour when starting terminal on wrong remote URL
         self.assertTrue(
@@ -141,14 +141,14 @@ def test_session_shell_vector(self, log_captured):
         self._assert_exec(':show shell_sh.vector', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'BOGUS'), log_captured)
 
         # Vectorlist methods ignore bogus vectors and just keep trying.
-        # TODO: should warn about unexistant vector, but seems too messy to fix
-        self._assert_exec('echo 1', '1', log_captured)
-        self._assert_exec(':show shell_sh.vector', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'system'), log_captured)
-
-        self._assert_exec(':set shell_sh.vector passthru', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'passthru'), log_captured)
-        self._assert_exec(':show shell_sh.vector', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'passthru'), log_captured)        
-        self._assert_exec('echo 1', '1', log_captured)
-        self._assert_exec(':show shell_sh.vector', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'passthru'), log_captured)
+        # TODO: warn about unexistant vector when implementing args completion
+        # self._assert_exec('echo 1', '1', log_captured)
+        # self._assert_exec(':show shell_sh.vector', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'system'), log_captured)
+        #
+        # self._assert_exec(':set shell_sh.vector passthru', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'passthru'), log_captured)
+        # self._assert_exec(':show shell_sh.vector', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'passthru'), log_captured)
+        # self._assert_exec('echo 1', '1', log_captured)
+        # self._assert_exec(':show shell_sh.vector', messages.sessions.set_module_s_s_s % ('shell_sh', 'vector', 'passthru'), log_captured)
 
 
     @log_capture()

From f1546505ac5f89127a365a64348fffbf162093aa Mon Sep 17 00:00:00 2001
From: Lucien A <lu.aubert84@gmail.com>
Date: Tue, 20 Jun 2023 23:34:28 +0200
Subject: [PATCH 4/4] clickify launcher

---
 core/modules.py            |  42 ++++++++++-
 core/terminal.py           |  29 ++++---
 requirements.txt           |  11 ++-
 tests/docker/entrypoint.sh |   2 +-
 weevely.py                 | 151 ++++++++++++++++++-------------------
 5 files changed, 140 insertions(+), 95 deletions(-)

diff --git a/core/modules.py b/core/modules.py
index 3f441e74..9c2cde8c 100644
--- a/core/modules.py
+++ b/core/modules.py
@@ -1,9 +1,49 @@
-from core import config
 import glob
 import os
 
+import click
+
+from core import config
+
 loaded = {}
 loaded_tree = {}
+plugin_folder = os.path.join(os.path.dirname(__file__), 'modules')
+
+
+class Manager(click.Group):
+
+    def __init__(self, **attrs):
+        attrs['invoke_without_command'] = True
+        super().__init__(**attrs)
+        self.help = 'BLABLABLA'
+
+    def list_commands(self, ctx):
+        rv = []
+        for filename in os.listdir(plugin_folder):
+            if filename.endswith('.py') and filename != '__init__.py':
+                rv.append(filename[:-3])
+        rv.sort()
+        return rv
+
+    def get_command(self, ctx, name):
+        ns = {}
+        fn = os.path.join(plugin_folder, name + '.py')
+        try:
+            with open(fn) as f:
+                code = compile(f.read(), fn, 'exec')
+                eval(code, ns, ns)
+        except FileNotFoundError:
+            return
+        return ns['cli']
+
+    def run(self, name, ctx=None, **kwargs):
+        if not ctx:
+            ctx = click.get_current_context()
+        cmd = self.get_command(ctx, name)
+
+        if not cmd:
+            return False  # @TODO raise instead ?
+        return ctx.forward(cmd, **kwargs)
 
 def load_modules(session):
     """ Load all modules """
diff --git a/core/terminal.py b/core/terminal.py
index ecf99ff1..599d2013 100644
--- a/core/terminal.py
+++ b/core/terminal.py
@@ -1,5 +1,6 @@
 import shlex
 import string
+import sys
 from urllib.parse import urlparse
 
 from mako import template
@@ -50,28 +51,36 @@ class Terminal:
 
     def __init__(self, session):
         self.session = session
-        self.prompt_session = PromptSession(history=FileHistory(config.history_path))
         self.completer = CustomCompleter(self)
 
         self._load_modules()
 
     def cmdloop(self):
+        prompt_session = PromptSession(
+            message=self.get_prompt_message,
+            history=FileHistory(config.history_path),
+            color_depth=ColorDepth.TRUE_COLOR,
+            complete_while_typing=True,
+            reserve_space_for_menu=10,
+            completer=self.completer,
+            key_bindings=self.kb,
+            style=style.default_style,
+        )
         self._print_intro()
+
+        if sys.stdin.isatty():
+            prompt = prompt_session.prompt
+        else:
+            prompt = sys.stdin.readline
+
         while True:
             try:
-                line = self.prompt_session.prompt(self.get_prompt_message,
-                                                  color_depth=ColorDepth.TRUE_COLOR,
-                                                  complete_while_typing=True,
-                                                  reserve_space_for_menu=10,
-                                                  completer=self.completer,
-                                                  key_bindings=self.kb,
-                                                  style=style.default_style,
-                                                  )
+                line = prompt()
                 line = self.precmd(line)
                 self.onecmd(line)
             except KeyboardInterrupt:
                 # Quit when pressing Ctrl-C while prompt is empty
-                if len(self.prompt_session.default_buffer.text) == 0:
+                if len(prompt_session.default_buffer.text) == 0:
                     raise EOFError
 
     def precmd(self, line):
diff --git a/requirements.txt b/requirements.txt
index f66d6d26..e14d5b5f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,10 @@
-prettytable
-Mako
-PyYAML
-python-dateutil
+prettytable~=3.7.0
+Mako~=1.2.4
+PyYAML~=6.0
+python-dateutil~=2.8.2
 PySocks
 pyOpenSSL
 prompt_toolkit
+click~=8.1.3
+testfixtures~=7.1.0
+Pygments~=2.15.1
\ No newline at end of file
diff --git a/tests/docker/entrypoint.sh b/tests/docker/entrypoint.sh
index 1d312081..ef48d17e 100755
--- a/tests/docker/entrypoint.sh
+++ b/tests/docker/entrypoint.sh
@@ -9,7 +9,7 @@ PWD="`python -c 'from tests import config;print(config.password)'`"
 # Generic environment setting install
 mkdir -p "$BASE_FOLDER"
 find -type f -name '*.pyc' -exec rm -f {} \;
-python ./weevely.py generate -obfuscator obfusc1_php "$PWD" "$AGENT"
+python ./weevely.py generate -o obfusc1_php "$PWD" "$AGENT"
 python ./weevely.py generate "$PWD" "$BASE_FOLDER"agent.phar
 
 a2enmod actions fcgid alias proxy_fcgi
diff --git a/weevely.py b/weevely.py
index 78a7270b..3947789d 100755
--- a/weevely.py
+++ b/weevely.py
@@ -3,109 +3,102 @@
 import os
 import pprint
 import sys
+import typing as t
 
-from core import generate
-from core import messages
-from core import modules
-from core.argparsers import CliParser
+import click
+
+from core import generate as _generate, messages, modules
 from core.config import agent_templates_folder_path, obfuscators_templates_folder_path
 from core.loggers import log, dlog
 from core.sessions import SessionURL, SessionFile
 from core.terminal import Terminal
-from core.weexceptions import FatalException
 
 if sys.stdout.encoding is None:
     print("Please set PYTHONIOENCODING=UTF-8 running 'export PYTHONIOENCODING=UTF-8' before starting Weevely.")
     exit(1)
 
 
-def main(arguments):
-
-    if arguments.command == 'generate':
+def list_templates(path):
+    return [
+        os.path.split(agent)[1].split('.')[0] for agent in
+        glob.glob('%s/*.tpl' % path)
+    ]
 
-        obfuscated = generate.generate(
-            password = arguments.password,
-            obfuscator = arguments.obfuscator,
-            agent = arguments.agent
-        )
 
-        generate.save_generated(obfuscated, arguments.path)
+agents_available = list_templates(agent_templates_folder_path)
+obfuscators_available = list_templates(obfuscators_templates_folder_path)
 
-        if arguments.path != '-':
-            log.info(messages.generate.generated_backdoor_with_password_s_in_s_size_i %
-                     (arguments.path,
-                      arguments.password,
-                      len(obfuscated)
-                      )
-                     )
 
-        return
+class DefaultGroup(click.Group):
+    def __init__(self, **attrs: t.Any):
+        attrs['invoke_without_command'] = True
+        super().__init__(**attrs)
 
-    elif arguments.command == 'terminal':
-        session = SessionURL(
-            url = arguments.url,
-            password = arguments.password
-        )
+    def parse_args(self, ctx, args):
+        if len(args) and args[0] not in ['terminal', 'generate', 'session']:
+            args.insert(0, 'terminal')
+        return super(DefaultGroup, self).parse_args(ctx, args)
 
-    elif arguments.command == 'session':
-        session = SessionFile(arguments.path)
 
+def run_cmd(_session, cmd):
     dlog.debug(
-        pprint.pformat(session)
+        pprint.pformat(_session)
     )
 
-    modules.load_modules(session)
+    modules.load_modules(_session)
 
-    if not arguments.cmd:
-        Terminal(session).cmdloop()
+    if not cmd:
+        Terminal(_session).cmdloop()
     else:
-        Terminal(session).onecmd(arguments.cmd)
-
-if __name__ == '__main__':
-
-    parser = CliParser(prog='weevely')
-    subparsers = parser.add_subparsers(dest = 'command')
-
-    terminalparser = subparsers.add_parser('terminal', help='Run terminal or command on the target')
-    terminalparser.add_argument('url', help = 'The agent URL')
-    terminalparser.add_argument('password', help = 'The agent password')
-    terminalparser.add_argument('cmd', help = 'Command', nargs='?')
+        Terminal(_session).onecmd(cmd)
+
+
+@click.group(cls=DefaultGroup)
+@click.pass_context
+def cli(ctx):
+    if ctx.invoked_subcommand is None:
+        click.echo(terminal.get_help(ctx))
+
+
+@cli.command()
+@click.argument('url')
+@click.argument('password')
+@click.argument('cmd', nargs=-1, required=False)
+def terminal(url, password, cmd):
+    """Connect to a Weevely agent at URL using PASSWORD."""
+    run_cmd(SessionURL(
+        url=url,
+        password=password
+    ), cmd)
+
+
+@cli.command()
+@click.argument('path')
+@click.argument('cmd', nargs=-1, required=False)
+def session(path, cmd):
+    run_cmd(SessionFile(path), cmd)
+
+
+@cli.command()
+@click.argument('password', required=True)
+@click.argument('path')
+@click.option('-o', '--obfuscator', default='phar', type=click.Choice(obfuscators_available))
+@click.option('-a', '--agent', default='obfpost_php', type=click.Choice(agents_available))
+def generate(password, path, obfuscator, agent):
+    """Generate a new agent at PATH using PASSWORD."""
+    obfuscated = _generate.generate(
+        password=password,
+        obfuscator=obfuscator,
+        agent=agent
+    )
 
-    sessionparser = subparsers.add_parser('session', help='Recover an existing session')
-    sessionparser.add_argument('path', help = 'Session file path')
-    sessionparser.add_argument('cmd', help = 'Command', nargs='?')
+    _generate.save_generated(obfuscated, path)
 
-    agents_available = [
-        os.path.split(agent)[1].split('.')[0] for agent in
-        glob.glob('%s/*.tpl' % agent_templates_folder_path)
-    ]
+    if path != '-':
+        log.info(
+            messages.generate.generated_backdoor_with_password_s_in_s_size_i %
+            (path, password, len(obfuscated))
+            )
 
-    obfuscators_available = [
-        os.path.split(agent)[1].split('.')[0] for agent in
-        glob.glob('%s/*.tpl' % obfuscators_templates_folder_path)
-    ]
 
-    generateparser = subparsers.add_parser('generate', help='Generate new agent')
-    generateparser.add_argument('password', help = 'Agent password')
-    generateparser.add_argument('path', help = 'Agent file path')
-    generateparser.add_argument(
-        '-obfuscator', #The obfuscation method
-        choices = obfuscators_available,
-        default = 'phar'
-        )
-    generateparser.add_argument(
-        '-agent', #The agent channel type
-        choices = agents_available,
-        default = 'obfpost_php'
-        )
-
-    parser.set_default_subparser('terminal')
-
-    arguments = parser.parse_args()
-
-    try:
-        main(arguments)
-    except (KeyboardInterrupt, EOFError):
-        log.info('Exiting.')
-    except FatalException as e:
-        log.critical('Exiting: %s' % e)
+cli()