From ab80be8a3d03d6afe39eacc78becfad515fe174d Mon Sep 17 00:00:00 2001
From: Dobin <dobin.rutishauser@gmail.com>
Date: Mon, 19 Feb 2024 19:50:35 +0000
Subject: [PATCH] ui: nicer log output and error handling

---
 helper.py          | 16 +++++++++++++---
 phases/injector.py |  3 ++-
 supermega.py       |  6 +++---
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/helper.py b/helper.py
index 0d6d7d7..3ea4573 100644
--- a/helper.py
+++ b/helper.py
@@ -35,8 +35,17 @@ def clean_files():
 
 
 def run_process_checkret(args, check=True):
-    ret = subprocess.run(args, 
-        capture_output=True)
+    ret = subprocess.CompletedProcess("", 666)
+    try:
+        ret = subprocess.run(args, capture_output=True)
+    except KeyboardInterrupt:
+        logger.warn("Caught KeyboardInterrupt, exiting gracefully...")
+    except subprocess.CalledProcessError as e:
+        logger.warn(f"Command '{e.cmd}' returned non-zero exit status {e.returncode}.")
+        # Handle the error case
+    except Exception as e:
+        logger.warn(f"An error occurred: {e}")
+        # Handle other exceptions
     
     with open("logs/cmdoutput.log", "ab") as f:
         cmd = "------------------------------------\n"
@@ -92,12 +101,13 @@ def delete_all_files_in_directory(directory_path):
 
 
 def rbrunmode_str(rbrunmode):
+    rbrunmode = str(rbrunmode)
     if rbrunmode == "1":
         return "change AddressOfEntryPoint"
     elif rbrunmode == "2":
         return "hijack branching instruction at Original Entry Point (jmp, call, ...)"
     else:
-        return "Invalid"
+        return "Invalid: {}".format(rbrunmode)
 
 
 def hexdump(data, addr = 0, num = 0):
diff --git a/phases/injector.py b/phases/injector.py
index 53844ad..e5e333d 100644
--- a/phases/injector.py
+++ b/phases/injector.py
@@ -8,7 +8,7 @@ import tempfile
 from pehelper import *
 from model import *
 from observer import observer
-
+from helper import rbrunmode_str
 from derbackdoorer.derbackdoorer import PeBackdoor
 
 logger = logging.getLogger("Injector")
@@ -23,6 +23,7 @@ def inject_exe(
     logger.info("--[ Injecting: {} into: {} -> {} mode {}".format(
         shellcode_in, exe_in, exe_out, inject_mode
     ))
+    logger.warn("--[ Inject mode: {}".format(rbrunmode_str(inject_mode)))
 
 
 
diff --git a/supermega.py b/supermega.py
index b6b7424..60d401b 100644
--- a/supermega.py
+++ b/supermega.py
@@ -229,7 +229,7 @@ def start(project: Project):
             exit_code = phases.injector.verify_injected_exe(project.inject_exe_out)
 
         elif project.try_start_final_infected_exe:
-            logger.info("--[ Start infected exe")
+            logger.info("--[ Start infected exe: {}".format(project.inject_exe_out))
             run_process_checkret([
                 project.inject_exe_out,
             ], check=False)
@@ -339,10 +339,10 @@ class ListHandler(logging.Handler):
 
 def setup_logging():
     root_logger = logging.getLogger()
-    root_logger.setLevel(logging.DEBUG)
+    root_logger.setLevel(logging.INFO)
 
     ch = logging.StreamHandler()
-    ch.setLevel(logging.DEBUG)
+    ch.setLevel(logging.INFO)
     ch.setFormatter(CustomFormatter())
 
     list_handler = ListHandler(log_messages)