sploit: Rework logger

The log module is updated to support binary encodings, colors, and for
improved compatibility with Python's print() builtin.

Encoding semantics are switched up, since it seems like some of the more
interesting encoding modes (from a CTF perspective) actually use
bytes-like objects as their high-level form (that is, bytes are encoded
to another form, such as hex, then decoded back to the original form).
So the logged value is now passed to encode instead of decode, and only
if the object is of type 'bytes', as unicode strings are now considered
out-of-scope for this operation.  Additionally, the bytes wrapper (b'')
is no longer visible in the logged content.

For readability, several standard colors have been defined for use
within Sploit:

        - RED:    Errors
        - YELLOW: Warnings
        - GREEN:  Status messages / Startup messages
        - WHITE:  Target output
        - GRAY:   User output / Alt text

Logging functions now support an optional color option to select the
desired color, and have specific defaults based on who is invoking the
log (see below...)

Logging functions are now also fully compatible with the builtin print()
function.  This is because Sploit now replaces the standard print() with
a logging function within the user's script (which is done to maintain
additional consistency of messages displayed in the console).

Function ilog (internal log) has default values tuned for the library's
convenience: Text goes to stderr, and is presented as status messages
(green).

Function elog (external log) has default values tuned for the user: Text
goes to stdout, and is presented as alt text to distinguish it from data
read from the target.  Within the user context, 'print' refers to this
function.

Signed-off-by: Malfurious <m@lfurio.us>
Signed-off-by: dusoleil <howcansocksbereal@gmail.com>

3 files changed, 56 insertions, 28 deletions

diff --git a/tools/sploit/sploit/comm.py b/tools/sploit/sploit/comm.py
index 7d6cd8c..2786c45 100644
--- a/tools/sploit/sploit/comm.py
+++ b/tools/sploit/sploit/comm.py
@@ -4,7 +4,7 @@ import os
 import sys
 import select
 
-from sploit.log import log
+from sploit.log import *
 from sploit.until import bind
 
 class Comm:
@@ -18,20 +18,20 @@ class Comm:
         data = os.read(self.back.stdin.fileno(), size)
         if(data == b''):
             raise BrokenPipeError('Tried to read on broken pipe')
-        if self.logonread : log(data)
+        if self.logonread : ilog(data, file=sys.stdout, color=NORMAL)
         return data
 
     def readline(self):
         data = self.back.stdin.readline()
         if(data == b''):
             raise BrokenPipeError('Tried to read on broken pipe')
-        if self.logonread : log(data)
+        if self.logonread : ilog(data, file=sys.stdout, color=NORMAL)
         return data
 
     def readall(self):
         data = b''
         for line in self.back.stdin:
-            log(line)
+            ilog(line, file=sys.stdout, color=NORMAL)
             data += line
         return data
 
@@ -45,7 +45,7 @@ class Comm:
             if(pred(data)):
                 break
         self.logonread = l
-        if self.logonread : log(data)
+        if self.logonread : ilog(data, file=sys.stdout, color=NORMAL)
         return data
 
     def readlineuntil(self, pred, /, *args, **kwargs):
@@ -65,7 +65,7 @@ class Comm:
         self.write(data + b'\n')
 
     def interact(self):
-        print("<--Interact Mode-->")
+        ilog("<--Interact Mode-->")
         stdin = sys.stdin.buffer
         os.set_blocking(self.back.stdin.fileno(), False)
         os.set_blocking(stdin.fileno(), False)
@@ -79,9 +79,11 @@ class Comm:
                 if(data == b''):
                     break
                 write(data)
+        def writeinput(write):
+            ilog(write, file=sys.stdout, color=NORMAL)
         readtable = {
                 stdin.fileno() : lambda : readall(stdin.readline, self.write),
-                self.back.stdin.fileno() : lambda : readall(self.back.stdin.readline, log)
+                self.back.stdin.fileno() : lambda : readall(self.back.stdin.readline, writeinput)
         }
         readtable[self.back.stdin.fileno()]()
         while(not brk):
@@ -97,17 +99,17 @@ class Comm:
                 break
         os.set_blocking(self.back.stdin.fileno(), True)
         os.set_blocking(stdin.fileno(), True)
-        print("<--Interact Mode Done-->")
+        ilog("<--Interact Mode Done-->")
 
 class Process:
     def __init__(self, args):
-        print(f"Running: {' '.join(args)}")
+        ilog(f"Running: {' '.join(args)}")
         self.proc = subprocess.Popen(args,
                 stdin=subprocess.PIPE,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT,
                 preexec_fn=lambda : os.setpgrp())
-        print(f"PID: {self.proc.pid}")
+        ilog(f"PID: {self.proc.pid}")
         self.stdin = self.proc.stdout
         self.stdout = self.proc.stdin
 
@@ -116,8 +118,8 @@ class Process:
         if(self.proc.poll() != None):
             return
         try:
-            print("Waiting on Target Program to End...")
-            print("Press Ctrl+C to Forcefully Kill It...")
+            ilog("Waiting on Target Program to End...")
+            ilog("Press Ctrl+C to Forcefully Kill It...")
             self.proc.wait()
         except KeyboardInterrupt:
             self.proc.kill()
@@ -135,11 +137,11 @@ class Pipes:
         self.pathout = os.path.join(dirname, "out")
         os.mkfifo(self.pathin)
         os.mkfifo(self.pathout)
-        print("Waiting on Target to Connect...")
-        print("<"+self.pathin+" >"+self.pathout)
+        ilog("Waiting on Target to Connect...", file=sys.stdout)
+        ilog(f"<{self.pathin} >{self.pathout}", file=sys.stdout)
         self.stdout = open(self.pathin, "wb")
         self.stdin = open(self.pathout, "rb")
-        print("Connected!")
+        ilog("Connected!")
 
     def __del__(self):
         try:
@@ -149,4 +151,3 @@ class Pipes:
             pass
         if getattr(self,'pathin',None) and os.path.exists(self.pathin) : os.unlink(self.pathin)
         if getattr(self,'pathout',None) and os.path.exists(self.pathout) : os.unlink(self.pathout)
-
diff --git a/tools/sploit/sploit/log.py b/tools/sploit/sploit/log.py
index cd9c3be..823b252 100644
--- a/tools/sploit/sploit/log.py
+++ b/tools/sploit/sploit/log.py
@@ -1,6 +1,32 @@
-ENCODING = ''
-def log(s):
-    if ENCODING != '':
-        s = s.decode(ENCODING)
-    print(s)
+import codecs
+import sys
 
+# https://docs.python.org/3/library/codecs.html#standard-encodings
+ENCODING = None
+
+ERROR   = 31
+WARNING = 33
+STATUS  = 32
+NORMAL  =  0
+ALT     = 90
+
+def enc_value(value, enc):
+    if type(value) is bytes:
+        if enc is not None:
+            value = codecs.encode(value, enc)
+        elif ENCODING is not None:
+            value = codecs.encode(value, ENCODING)
+        value = str(value)[2:-1] # strip b''
+    return str(value)
+
+def generic_log(*values, sep, end, file, flush, enc, color):
+    string = sep.join([ enc_value(x, enc) for x in values ])
+    print(f'\033[{color}m{string}\033[0m', end=end, file=file, flush=flush)
+
+# For library internal use
+def ilog(*values, sep=' ', end='\n', file=sys.stderr, flush=True, enc=None, color=STATUS):
+    generic_log(*values, sep=sep, end=end, file=file, flush=flush, enc=enc, color=color)
+
+# For external use in user script (via print = elog)
+def elog(*values, sep=' ', end='\n', file=sys.stdout, flush=True, enc=None, color=ALT):
+    generic_log(*values, sep=sep, end=end, file=file, flush=flush, enc=enc, color=color)
diff --git a/tools/sploit/sploit/main.py b/tools/sploit/sploit/main.py
index 8456029..0a34429 100644
--- a/tools/sploit/sploit/main.py
+++ b/tools/sploit/sploit/main.py
@@ -4,6 +4,7 @@ import tempfile
 import traceback
 
 from sploit.comm import *
+from sploit.log import *
 
 def main():
     parser = ArgumentParser(description='Execute Sploit script against target')
@@ -17,7 +18,7 @@ def main():
         pipe(args.script)
 
 def pipe(script):
-    print("Running in Pipe Mode...")
+    ilog("Running in Pipe Mode...")
     with tempfile.TemporaryDirectory() as tmpdir:
         while(True):
             try:
@@ -28,21 +29,21 @@ def pipe(script):
             del p
 
 def target(script, target):
-    print("Running in Target Mode...")
+    ilog("Running in Target Mode...")
     runscript(script, Comm(Process(target)))
 
 def runscript(script, comm):
     try:
-        print("Running Script...")
+        ilog("Running Script...")
         code = compile(open(script).read(), script, 'exec')
-        exec(code, {'io': comm})
-        print("Script Finished!")
+        exec(code, {'io': comm, 'print': elog})
+        ilog("Script Finished!")
         comm.readall()
         return
     except KeyboardInterrupt:
         pass
     except:
-        traceback.print_exc()
+        ilog(traceback.format_exc(), end='', color=ERROR)
     finally:
         gc.collect()
-    print("Script Ended Early!")
+    ilog("Script Ended Early!", color=WARNING)