refactor: cleanup, error handling

observer.py

@@ -1,11 +1,6 @@
-import json
-import pprint
-from capstone import Cs, CS_ARCH_X86, CS_MODE_64
 from typing import List, Dict
 
 from pe.r2helper import r2_disas
-from utils import delete_all_files_in_directory
-from model.defs import *
 
 
 class Observer():
@@ -51,32 +46,6 @@ class Observer():
         #self.write_to_file(name + ".disas.ascii", ret['color'])
         #self.write_to_file(name + ".hex", ret['hexdump'])
         #self.write_to_file_bin(name + ".bin", data)
-        #self.idx += 1
-
-  
-    #def write_to_file(self, filename, data):
-    #    if not self.active:
-    #        return
-    #    with open("{}/{}-{}".format(logs_dir, self.idx, filename), "w") as f:
-    #        f.write(data)
-
-
-    #def write_to_file_bin(self, filename, data):
-    #    if not self.active:
-    #        return
-    #    with open("{}/{}-{}".format(logs_dir, self.idx, filename), "wb") as f:
-    #        f.write(data)
-
-
-    #def clean_files(self):
-    #    delete_all_files_in_directory(f"{logs_dir}/")
-    #    self.idx = 0
-    #    self.logs = []
-
-
-    #def __str__(self):
-    #    s = "<todo>"
-    #    return s
-
+ 
 
 observer = Observer()

supermega.py

@@ -101,7 +101,9 @@ def main():
     exit(exit_code)
 
 
-def start(settings: Settings):
+def start(settings: Settings) -> int:
+    """Main entry point for the application. Will handle log files and cleanup"""
+
     # Delete: all old files
     if settings.cleanup_files_on_start:
         clean_files()
@@ -121,25 +123,7 @@ def start(settings: Settings):
         clean_files()
 
     write_logs()
-
-
-def write_logs():
-    # Our log output
-    with open(f"{logs_dir}/supermega.log", "w") as f:
-        for line in observer.get_logs():
-            f.write(line + "\n")
-
-    # Stdout of executed commands
-    with open(f"{logs_dir}/cmdoutput.log", "w") as f:
-        for line in observer.get_cmd_output():
-            f.write(line)
-
-    # Write all files
-    idx = 0
-    for name, data in observer.files:
-        with open(f"{logs_dir}/{idx}-{name}", "w") as f:
-            f.write(data)
-        idx += 1
+    return 0
 
 
 def start_real(settings: Settings):
@@ -216,12 +200,33 @@ def start_real(settings: Settings):
         if settings.verify:
             logger.info("--[ Verify infected exe")
             payload_exit_code = phases.injector.verify_injected_exe(settings.inject_exe_out)
-            logging.info("Payload xit code: {}".format(payload_exit_code))
+            logging.info("Payload exit code: {}".format(payload_exit_code))
+            if payload_exit_code != 0:
+                raise Exception("Payload exit code: {}".format(payload_exit_code))
         elif settings.try_start_final_infected_exe:
             logger.info("--[ Start infected exe: {}".format(settings.inject_exe_out))
             run_process_checkret([
                 settings.inject_exe_out,
-            ], check=False)
+            ], check=True)
+
+
+def write_logs():
+    # Our log output
+    with open(f"{logs_dir}/supermega.log", "w") as f:
+        for line in observer.get_logs():
+            f.write(line + "\n")
+
+    # Stdout of executed commands
+    with open(f"{logs_dir}/cmdoutput.log", "w") as f:
+        for line in observer.get_cmd_output():
+            f.write(line)
+
+    # Write all files
+    idx = 0
+    for name, data in observer.files:
+        with open(f"{logs_dir}/{idx}-{name}", "w") as f:
+            f.write(data)
+        idx += 1
 
 
 def obfuscate_shc_loader(file_shc_in, file_shc_out):