clickify launcher

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 """

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):

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

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

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()