From a0fd2ecc1e8d2bde8b218e75b65cedc641b23cff Mon Sep 17 00:00:00 2001 From: Dobin Date: Fri, 16 Feb 2024 09:13:28 +0000 Subject: [PATCH] refactor: templater --- defs.py | 12 +++++++++++- phases/assembler.py | 5 ----- phases/templater.py | 41 +++++++++++++++++++++++------------------ supermega.py | 13 +++++++------ 4 files changed, 41 insertions(+), 30 deletions(-) diff --git a/defs.py b/defs.py index 483993f..f40bc80 100644 --- a/defs.py +++ b/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") \ No newline at end of file diff --git a/phases/assembler.py b/phases/assembler.py index 4e2b075..bb73fb7 100644 --- a/phases/assembler.py +++ b/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") diff --git a/phases/templater.py b/phases/templater.py index 1babf07..bb00824 100644 --- a/phases/templater.py +++ b/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") \ No newline at end of file + shutil.copy("source/iat_reuse/main.c", main_c_file) \ No newline at end of file diff --git a/supermega.py b/supermega.py index 1900d6e..5074817 100644 --- a/supermega.py +++ b/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: