diff --git a/app/storage.py b/app/storage.py index 036dd37..40dd94e 100644 --- a/app/storage.py +++ b/app/storage.py @@ -1,4 +1,6 @@ import pickle +import os +import yaml from typing import List, Tuple from model.settings import Settings @@ -14,20 +16,31 @@ class Storage(): def __init__(self): self.data: List[Project] = self.get_data() - def get_project(self, name): + def get_project(self, name: str) -> Project: for project in self.data: if project.name == name: return project return None - def add_project(self, project): + def add_project(self, project: Project): + # data self.data.append(project) self.save_data() - def get_data(self): + # directories and contents + os.makedirs("app/projects/{}".format(project.name), exist_ok=True) + with open("app/projects/{}/settings.yaml".format(project.name), "w") as f: + f.write(yaml.dump(project.settings)) + + def get_data(self) -> List[Project]: + # if file does not exist, create an empty one + if not os.path.exists("app/data.pickle"): + with open("app/data.pickle", "wb") as f: + f.write(pickle.dumps([])) + with open("app/data.pickle", "rb") as f: - data = f.read() - data = pickle.loads(data) + data_raw = f.read() + data: List[Project] = pickle.loads(data_raw) return data def save_data(self): diff --git a/app/templates/index.html b/app/templates/index.html index 442d97a..8b45358 100644 --- a/app/templates/index.html +++ b/app/templates/index.html @@ -17,7 +17,7 @@ {% endfor %} - Add + Add Project \ No newline at end of file diff --git a/app/templates/project.html b/app/templates/project.html index 08a9110..2bd6081 100644 --- a/app/templates/project.html +++ b/app/templates/project.html @@ -14,7 +14,7 @@
- -
+ + -
- - -
-
diff --git a/app/views.py b/app/views.py index bd1c8ce..3c92cd6 100644 --- a/app/views.py +++ b/app/views.py @@ -9,6 +9,7 @@ from pygments.lexers import CLexer, NasmLexer, DiffLexer, HexdumpLexer from pygments.formatters import HtmlFormatter import difflib from ansi2html import Ansi2HTMLConverter +import shutil from config import config from model.settings import Settings @@ -26,7 +27,6 @@ thread_running = False @views.route("/") def index(): - print(storage.data) return render_template('index.html', data=storage.data) @@ -96,9 +96,11 @@ def add_project(): settings.inject_style = InjectStyle[inject_style] if storage.get_project(project_name) != None: + # overwrite project project = storage.get_project(project_name) project.settings = settings else: + # add new project project = Project(project_name, settings) project.settings = settings settings.project_name = project_name @@ -132,11 +134,18 @@ def add_project(): ) -def supermega_thread(settings: Settings): +def supermega_thread(settings: Settings, project_name: str): global thread_running start(settings) thread_running = False + # copy generated file to project folder + file_basename = os.path.basename(settings.inject_exe_out) + shutil.copy( + settings.inject_exe_out, + "app/projects/{}/{}".format(project_name, file_basename) + ) + @views.route("/start_project", methods=['POST', 'GET']) def start_project(): @@ -152,7 +161,7 @@ def start_project(): project = storage.get_project(project_name) project.settings.try_start_final_infected_exe = try_start - thread = Thread(target=supermega_thread, args=(project.settings, )) + thread = Thread(target=supermega_thread, args=(project.settings, project_name, )) thread.start() thread_running = True diff --git a/sender.py b/sender.py index 1ea33cd..5ab8ddb 100644 --- a/sender.py +++ b/sender.py @@ -8,6 +8,9 @@ import shutil from config import config +logger = logging.getLogger("Sender") + + def scannerDetectsBytes(data: bytes, filename: str, useBrotli=True, verify=False): params = { 'filename': filename, 'brotli': useBrotli, 'verify': verify } if useBrotli: @@ -16,11 +19,12 @@ def scannerDetectsBytes(data: bytes, filename: str, useBrotli=True, verify=False scanData = data timeStart = time.time() - print("Send to exec/exe: {}".format(params)) + logger.info("Send to exec/exe: {}".format(params)) res = req.post("{}/exec/exe".format(config.get("avred_server")), params=params, data=scanData, timeout=10) jsonRes = res.json() - print("Response: {}".format(jsonRes)) scanTime = round(time.time() - timeStart, 3) + logger.info("Response: {}s: {}".format(scanTime, jsonRes)) + # basically internal server error, e.g. AMSI not working if res.status_code != 200: diff --git a/supermega.py b/supermega.py index 155bcbf..c33379b 100644 --- a/supermega.py +++ b/supermega.py @@ -187,10 +187,14 @@ def start(settings: Settings): observer.add_code("exe_final", extract_code_from_exe_file_ep(settings.inject_exe_out, 300)) if config.get("avred_server") != "": + filename = os.path.basename(settings.inject_exe_in) with open(settings.inject_exe_out, "rb") as f: data = f.read() - scannerDetectsBytes(data, "test.exe", useBrotli=True, verify=settings.verify) - + try: + scannerDetectsBytes(data, filename, useBrotli=True, verify=settings.verify) + except Exception as e: + logger.error(f'Error scanning: {e}') + return 4 else: # Start/verify it at the end if settings.verify: