From 4410685d0d162468ea318275634953782175892a Mon Sep 17 00:00:00 2001 From: Dobin Date: Fri, 29 Mar 2024 20:29:33 +0000 Subject: [PATCH] refactor: unify project handling (cmdline = projects/default) --- app/storage.py | 9 +-------- app/templates/project.html | 2 +- app/views.py | 25 ++----------------------- helper.py | 12 ++++++++++++ model/project.py | 29 ++++++++++++++++++++++++++++- model/settings.py | 12 +++--------- supermega.py | 8 ++++++-- 7 files changed, 53 insertions(+), 44 deletions(-) diff --git a/app/storage.py b/app/storage.py index a9e037e..ce2961d 100644 --- a/app/storage.py +++ b/app/storage.py @@ -4,15 +4,8 @@ import yaml import pickle from typing import List, Tuple -from model.settings import Settings from model.defs import * - - -class WebProject(): - def __init__(self, name: str, settings: Settings): - self.name = name - self.settings: Settings = settings - self.comment: str = "" +from model.project import WebProject class Storage(): diff --git a/app/templates/project.html b/app/templates/project.html index 42af54e..71bf6af 100644 --- a/app/templates/project.html +++ b/app/templates/project.html @@ -50,7 +50,7 @@ diff --git a/app/views.py b/app/views.py index 016453c..6b9c24e 100644 --- a/app/views.py +++ b/app/views.py @@ -9,7 +9,6 @@ from pygments.lexers import CLexer, NasmLexer, DiffLexer, HexdumpLexer from pygments.formatters import HtmlFormatter import difflib from ansi2html import Ansi2HTMLConverter -import shutil import subprocess import time from datetime import datetime @@ -25,6 +24,7 @@ from phases.injector import verify_injected_exe from phases.compiler import compile_dev from phases.assembler import asm_to_shellcode from helper import run_process_checkret +from model.project import prepare_project views = Blueprint('views', __name__) @@ -184,12 +184,6 @@ def add_project(): decoder_style = request.form['decoder_style'] settings.decoder_style = DecoderStyle[decoder_style] - exec_style = request.form['exec_style'] - settings.exec_style = ExecStyle[exec_style] - - inject_style = request.form['inject_style'] - settings.inject_style = InjectStyle[inject_style] - if storage.get_project(project_name) != None: # overwrite project project = storage.get_project(project_name) @@ -242,22 +236,7 @@ def build_project(project_name): project = storage.get_project(project_name) project.settings.try_start_final_infected_exe = False - - src = "{}{}/".format(PATH_CARRIER, project.settings.source_style.value) - dst = "{}{}/".format(PATH_WEB_PROJECT, project_name) - - # delete all files in dst directory - for file in os.listdir(dst): - if file == "project.pickle": - continue - os.remove(dst + file) - - # copy *.c *.h files from src directory to dst directory - for file in os.listdir(src): - if file.endswith(".c") or file.endswith(".h"): - logger.info("Copy {} to {}".format(src + file, dst)) - shutil.copy2(src + file, dst) - + prepare_project(project_name, project.settings) thread = Thread(target=supermega_thread, args=(project.settings, )) thread.start() thread_running = True diff --git a/helper.py b/helper.py index e17c66e..9ee8972 100644 --- a/helper.py +++ b/helper.py @@ -3,7 +3,9 @@ import os import pathlib import glob import logging +import pickle +from model.project import WebProject from config import config from model.defs import * from observer import observer @@ -13,6 +15,15 @@ logger = logging.getLogger("Helper") SHC_VERIFY_SLEEP = 0.1 +def write_webproject(project_name, settings): + filepath = "{}project.pickle".format(settings.main_dir) + logger.info("Write project to: {}".format(filepath)) + webProject = WebProject(project_name, settings) + webProject.comment = "Created by command line interface" + with open(filepath, "wb") as f: + pickle.dump(webProject, f) + + def clean_tmp_files(): files_to_clean = [ # compile artefacts in current working dir @@ -23,6 +34,7 @@ def clean_tmp_files(): for file in files_to_clean: pathlib.Path(file).unlink(missing_ok=True) + def clean_files(settings): logger.info("--( Remove old files") diff --git a/model/project.py b/model/project.py index d9c4276..bd9c377 100644 --- a/model/project.py +++ b/model/project.py @@ -1,14 +1,22 @@ import logging +import shutil +from model.defs import * from model.payload import Payload from model.exehost import ExeHost from model.settings import Settings from model.carrier import Carrier - logger = logging.getLogger("Project") +class WebProject(): + def __init__(self, name: str, settings: Settings): + self.name = name + self.settings: Settings = settings + self.comment: str = "" + + class Project(): def __init__(self, settings: Settings): self.name: str = "" @@ -26,3 +34,22 @@ class Project(): self.payload.init() self.exe_host.init() self.carrier.init() + + +def prepare_project(project_name, settings): + src = "{}{}/".format(PATH_CARRIER, settings.source_style.value) + dst = "{}{}/".format(PATH_WEB_PROJECT, project_name) + + # delete all files in dst directory + for file in os.listdir(dst): + if file == "project.pickle": + continue + if file.startswith("."): + continue + os.remove(dst + file) + + # copy *.c *.h files from src directory to dst directory + for file in os.listdir(src): + if file.endswith(".c") or file.endswith(".h"): + logger.info("Copy {} to {}".format(src + file, dst)) + shutil.copy2(src + file, dst) \ No newline at end of file diff --git a/model/settings.py b/model/settings.py index d4d21aa..8b365a6 100644 --- a/model/settings.py +++ b/model/settings.py @@ -1,5 +1,8 @@ +import logging from model.defs import * +logger = logging.getLogger("Views") + class Settings(): def __init__(self, web=""): @@ -28,15 +31,6 @@ class Settings(): self.generate_shc_from_asm: bool = True - def prep(self): - self.main_dir = "{}{}/".format(PATH_CARRIER, self.source_style.value) - self.template_path = self.main_dir + "template.c" - self.main_c_path = self.main_dir + "main.c" - self.main_asm_path = self.main_dir + "main.asm" - self.main_exe_path = self.main_dir + "main.exe" - self.main_shc_path = self.main_dir + "main.bin" - - def prep_web(self, project_name): self.main_dir = "{}{}/".format(PATH_WEB_PROJECT, project_name) self.template_path = self.main_dir + "template.c" diff --git a/supermega.py b/supermega.py index 5d05969..4ad2b55 100644 --- a/supermega.py +++ b/supermega.py @@ -14,7 +14,7 @@ import phases.injector from observer import observer from pe.pehelper import extract_code_from_exe_file_ep from sender import scannerDetectsBytes -from model.project import Project +from model.project import Project, prepare_project from model.settings import Settings from model.defs import * from log import setup_logging @@ -96,7 +96,8 @@ def main(): settings.inject_exe_in = args.inject settings.inject_exe_out = args.inject.replace(".exe", ".infected.exe") - settings.prep() + settings.prep_web("default") + write_webproject("default", settings) exit_code = start(settings) exit(exit_code) @@ -112,6 +113,9 @@ def start(settings: Settings) -> int: # And logs observer.reset() + # Prepare the project: copy all files to projects// + prepare_project("default", settings) + # Do the thing and catch the errors try: start_real(settings)