feature: UI/UX improvements

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_raw = f.read()
-            data = pickle.loads(data)
+            data: List[Project] = pickle.loads(data_raw)
         return data
 
     def save_data(self):

app/templates/index.html

@@ -17,7 +17,7 @@
       {% endfor %}
     </ul>
 
-    <a href="/add_project">Add</a>
+    <a href="/add_project">Add Project</a>
   </div>
   </body>
 </html>

app/templates/project.html

@@ -14,7 +14,7 @@
 
         <div class="row">
             <div class="col-3">
-                <input type="text" name="comment" class="hidden form-control" value="{{project.comment}}"
+                <input type="text" name="comment" class="hidden form-control" placeholder="Comment" value="{{project.comment}}"
                 placeholder="" aria-label="PROJECTNAME" aria-describedby="basic-addon1">
             
                 <select class="form-select" name="shellcode" aria-label="SHELLCODE">

app/templates/project_add_get.html

@@ -14,10 +14,10 @@
 
     <form method="POST" enctype="multipart/form-data" action="/add_project">
 
-    <input type="text" name="project_name" class="form-control" placeholder="Projekt" aria-label="PROJECTNAME" aria-describedby="basic-addon1">
-
     <div class="row row-cols-lg-auto g-3">
 
+    <input type="text" name="project_name" class="form-control" placeholder="Projekt name" aria-label="PROJECTNAME" aria-describedby="basic-addon1">
+    
     <select class="form-select" name="shellcode" aria-label="SHELLCODE">
         {% for shellcode in shellcodes %}
         <option value="{{shellcode}}">{{shellcode}}</option>
@@ -60,13 +60,6 @@
         {% endfor %}
     </select>
 
-    <div class="form-check">
-        <input class="form-check-input" type="checkbox" value="" id="flexCheckDefault">
-        <label class="form-check-label" for="flexCheckDefault">
-          Execute after injection
-        </label>
-      </div>
-
     <button class="btn btn-primary" type="submit" value="save">Save</button>
 
     </div>

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
 

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:

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: