refactor: dedicated tester

log.py

@@ -52,16 +52,16 @@ def writelog():
         for line in log_messages:
             f.write(line + "\n")
 
-def setup_logging():
+def setup_logging(level = logging.INFO):
     root_logger = logging.getLogger()
-    root_logger.setLevel(logging.INFO)
+    root_logger.setLevel(level)
 
     ch = logging.StreamHandler()
-    ch.setLevel(logging.INFO)
+    ch.setLevel(level)
     ch.setFormatter(CustomFormatter())
 
     list_handler = ListHandler(log_messages)
-    list_handler.setLevel(logging.DEBUG)
+    list_handler.setLevel(level)
     list_handler.setFormatter(CustomFormatter())
 
     root_logger.addHandler(ch)

phases/injector.py

@@ -113,7 +113,7 @@ def injected_fix_data(superpe: SuperPe, carrier: Carrier, exe_host: ExeHost):
     
     rm = exe_host.get_rdata_relocmanager()
 
-    if True:
+    if True:  # FIXME this is a hack which is sometimes necessary
         sect_data_copy = peSection.pefile_section.get_data()
         string_off = find_first_utf16_string_offset(sect_data_copy)
         if string_off == None:

supermega.py

@@ -42,39 +42,12 @@ def main():
     parser.add_argument('--short-call-patching', action='store_true', help='Make short calls long. You will know when you need it.')
     parser.add_argument('--no-clean-at-start', action='store_true', help='Debug: Dont remove any temporary files at start')
     parser.add_argument('--no-clean-at-exit', action='store_true', help='Debug: Dont remove any temporary files at exit')
-    parser.add_argument('--verify', type=str, help='Debug: Perform verification: std/iat')
     parser.add_argument('--show', action='store_true', help='Debug: Show tool output')
     args = parser.parse_args()
 
     if args.show:
         config.ShowCommandOutput = True
 
-    if args.verify:
-        settings.payload_path = "data/shellcodes/createfile.bin"
-        settings.verify = True
-
-        settings.try_start_final_infected_exe = False
-
-        if args.verify == "peb":
-            settings.source_style = SourceStyle.peb_walk
-            settings.inject_mode = InjectStyle.BackdoorCallInstr
-            settings.inject_exe_in = "data/exes/7z.exe"
-            settings.inject_exe_out = "data/exes/7z-verify.exe"
-        elif args.verify == "iat":
-            settings.source_style = SourceStyle.iat_reuse
-            settings.inject_mode = InjectStyle.BackdoorCallInstr
-            settings.inject_exe_in = "data/exes/procexp64.exe"
-            settings.inject_exe_out = "data/exes/procexp64-verify.exe"
-        elif args.verify == "rwx":
-            settings.source_style = SourceStyle.peb_walk
-            settings.inject_mode = InjectStyle.ChangeEntryPoint  # ,2 is broken atm
-            settings.inject_exe_in = "data/exes/wifiinfoview.exe"
-            settings.inject_exe_out = "data/exes/wifiinfoview.exe-verify.exe"
-        else:
-            logger.info("Unknown verify option {}, use std/iat".format(args.verify))
-            return
-
-    else:
     settings.try_start_final_infected_exe = args.start_injected
     settings.cleanup_files_on_start = not args.no_clean_at_start
     settings.cleanup_files_on_exit =not args.no_clean_at_exit
@@ -211,10 +184,10 @@ def start(settings: Settings):
         phases.injector.inject_exe(main_shc_file, settings, project)
     except PermissionError as e:
         logger.error(f'Error writing file: {e}')
-        return exit(2)
+        return 2
     except Exception as e:
         logger.error(f'Error injecting: {e}')
-        return exit(3)
+        return 3
     
     observer.add_code("exe_final", extract_code_from_exe_file_ep(settings.inject_exe_out, 300))
 

tester.py

@@ -0,0 +1,60 @@
+from typing import Dict
+
+from helper import *
+from config import config
+
+from model.settings import Settings
+from log import setup_logging
+from supermega import start
+
+
+def main():
+    """Argument parsing for when called from command line"""
+    logger.info("Super Mega")
+    config.load()
+
+    settings = Settings()
+    settings.payload_path = "data/shellcodes/createfile.bin"
+    settings.verify = True
+    settings.try_start_final_infected_exe = False
+    
+    # 7z, peb-walk, change-entrypoint
+    settings.source_style = SourceStyle.peb_walk
+    settings.inject_mode = InjectStyle.ChangeEntryPoint
+    settings.inject_exe_in = "data/exes/7z.exe"
+    settings.inject_exe_out = "data/exes/7z-verify.exe"
+    if start(settings) != 0:
+        print("Error")
+        return 1
+
+    # 7z, peb-walk, hijack
+    settings.source_style = SourceStyle.peb_walk
+    settings.inject_mode = InjectStyle.BackdoorCallInstr
+    settings.inject_exe_in = "data/exes/7z.exe"
+    settings.inject_exe_out = "data/exes/7z-verify.exe"
+    if start(settings) != 0:
+        print("Error")
+        return 1
+
+    # procexp, iat-reuse, change-entrypoint
+    settings.source_style = SourceStyle.iat_reuse
+    settings.inject_mode = InjectStyle.ChangeEntryPoint
+    settings.inject_exe_in = "data/exes/procexp64.exe"
+    settings.inject_exe_out = "data/exes/procexp64-verify.exe"
+    if start(settings) != 0:
+        print("Error")
+        return 1
+
+    # procexp, iat-reuse, change-entrypoint
+    settings.source_style = SourceStyle.iat_reuse
+    settings.inject_mode = InjectStyle.ChangeEntryPoint
+    settings.inject_exe_in = "data/exes/procexp64.exe"
+    settings.inject_exe_out = "data/exes/procexp64-verify.exe"
+    if start(settings) != 0:
+        print("Error")
+        return 1
+
+
+if __name__ == "__main__":
+    setup_logging(level=logging.WARN)
+    main()