refactor: templater

defs.py

@@ -1,10 +1,13 @@
 from enum import Enum
-
+import os
 
 class FilePath(str):
     pass
 
 
+# Correlated with real template files
+# in plugins/
+
 class AllocStyle(Enum):
     RWX = "rwx_1"
     #RW_X = "rw_x"
@@ -28,3 +31,10 @@ class SourceStyle(Enum):
     peb_walk = 1
     iat_reuse = 2
 
+
+build_dir = "build"
+
+main_c_file = os.path.join(build_dir, "main.c")
+main_asm_file = os.path.join(build_dir, "main.asm")
+main_exe_file = os.path.join(build_dir, "main.exe")
+main_shc_file = os.path.join(build_dir, "main.bin")

phases/assembler.py

@@ -1,13 +1,8 @@
-import pefile
-import pprint
 import logging
 
 from model import *
-from helper import *
 from config import config
 from observer import observer
-from project import project
-from helper import *
 from pehelper import *
 
 logger = logging.getLogger("Assembler")

phases/templater.py

@@ -4,14 +4,13 @@ import shutil
 import logging
 
 from helper import *
-from config import config
-from project import project
-from model import *
 from observer import observer
+from defs import *
 
 use_templates = True
 logger = logging.getLogger("Assembler")
 
+
 # INPUT:
 #   plugins/
 #   source/
@@ -19,26 +18,30 @@ logger = logging.getLogger("Assembler")
 # Output:
 #   build/main.c
 #   build/*.h
-def create_c_from_template():
+def create_c_from_template(
+    source_style: SourceStyle, 
+    alloc_style: AllocStyle, 
+    exec_style: ExecStyle,
+    decoder_style: DecoderStyle,
+    build_dir: FilePath,
+):
     plugin_allocator = ""
     plugin_decoder = ""
     plugin_executor = ""
 
-    with open("plugins/allocator/rwx_1.c", "r", encoding='utf-8') as file:
+    filepath = "plugins/allocator/{}.c".format(alloc_style.value)
+    with open(filepath, "r", encoding='utf-8') as file:
         plugin_allocator = file.read()
 
-    if project.decoder_style == DecoderStyle.PLAIN_1:
-        with open("plugins/decoder/plain_1.c", "r", encoding='utf-8') as file:
-            plugin_decoder = file.read()
-    elif project.decoder_style == DecoderStyle.XOR_1:
-        with open("plugins/decoder/xor_1.c", "r", encoding='utf-8') as file:
+    filepath = "plugins/decoder/{}.c".format(decoder_style.value)
+    with open(filepath, "r", encoding='utf-8') as file:
             plugin_decoder = file.read()
 
+    filepath = "plugins/executor/{}.c".format(exec_style.value)
     with open("plugins/executor/direct_1.c", "r", encoding='utf-8') as file:
         plugin_executor = file.read()
-
     
-    if project.source_style == SourceStyle.peb_walk:
+    if source_style == SourceStyle.peb_walk:
         if use_templates:
             with open("source/peb_walk/template.c", 'r', encoding='utf-8') as file:
                 template_content = file.read()
@@ -50,17 +53,19 @@ def create_c_from_template():
                 'plugin_decoder': plugin_decoder,
                 'plugin_executor': plugin_executor,
             })
-            with open("build/main.c", "w", encoding='utf-8') as file:
+            with open(main_c_file, "w", encoding='utf-8') as file:
                 file.write(rendered_template)
                 observer.add_text("main_c_rendered", rendered_template)
-            shutil.copy("source/peb_walk/peb_lookup.h", "build/peb_lookup.h")
 
+            # TODO PEB
+            shutil.copy("source/peb_walk/peb_lookup.h", "build/peb_lookup.h")
         else:
             observer.add_text("main_c", file_readall_text("source/peb_walk/main.c"))
-            shutil.copy("source/peb_walk/main.c", "build/main.c")
+            shutil.copy("source/peb_walk/main.c", main_c_file)
+            # TODO PEB
             shutil.copy("source/peb_walk/peb_lookup.h", "build/peb_lookup.h")
 
-    elif project.source_style == SourceStyle.iat_reuse:
+    elif source_style == SourceStyle.iat_reuse:
         if use_templates:
             with open("source/iat_reuse/template.c", 'r', encoding='utf-8') as file:
                 template_content = file.read()
@@ -71,9 +76,9 @@ def create_c_from_template():
                 'plugin_decoder': plugin_decoder,
                 'plugin_executor': plugin_executor,
             })
-            with open("build/main.c", "w", encoding='utf-8') as file:
+            with open(main_c_file, "w", encoding='utf-8') as file:
                 file.write(rendered_template)
                 observer.add_text("main_c_rendered", rendered_template)
         else:
             observer.add_text("main_c", file_readall_text("source/iat_reuse/main.c"))
-            shutil.copy("source/iat_reuse/main.c", "build/main.c")
+            shutil.copy("source/iat_reuse/main.c", main_c_file)

supermega.py

@@ -17,11 +17,6 @@ import phases.injector
 from observer import observer
 from project import project
 
-main_c_file = os.path.join(project.build_dir, "main.c")
-main_asm_file = os.path.join(project.build_dir, "main.asm")
-main_exe_file = os.path.join(project.build_dir, "main.exe")
-main_shc_file = os.path.join(project.build_dir, "main.bin")
-
 
 # ANSI escape sequences for colors
 class LogColors:
@@ -163,7 +158,13 @@ def start():
     logger.warning("--[ SourceStyle: {}".format(project.source_style.name))
 
     # Copy: loader C files into working directory: build/
-    phases.templater.create_c_from_template()
+    phases.templater.create_c_from_template(
+        source_style = project.source_style,
+        alloc_style  = project.alloc_style,
+        exec_style   = project.exec_style,
+        decoder_style= project.decoder_style,
+        build_dir    = project.build_dir,
+    )
 
     # Convert: C -> ASM
     if project.generate_asm_from_c: