Add a lil color

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

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\''
@@ -54,20 +42,20 @@ 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 = """
-[+] weevely ${version}
+    <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.
+  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>'

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