refactor: make log a class MyLog

log.py

@@ -2,7 +2,6 @@ import logging
 
 from model.defs import *
 
-log_messages = []
 
 
 # Logging
@@ -18,6 +17,7 @@ class LogColors:
     BOLD = '\033[1m'
     UNDERLINE = '\033[4m'
 
+
 class CustomFormatter(logging.Formatter):
     #format = "%(asctime)s - %(name)-12s - [%(levelname)-8s] - %(message)s (%(filename)s:%(lineno)d)"
     format = "(%(filename)-12s) %(message)s"
@@ -35,6 +35,7 @@ class CustomFormatter(logging.Formatter):
         formatter = logging.Formatter(log_fmt, datefmt="%Y-%m-%d %H:%M:%S")
         return formatter.format(record)
 
+
 class ListHandler(logging.Handler):
     def __init__(self, log_list):
         super().__init__()
@@ -46,33 +47,39 @@ class ListHandler(logging.Handler):
         self.log_list.append(log_entry)
 
 
-def writelog():
-    # write log to file
-    with open(f"{logs_dir}/supermega.log", "w") as f:
-        for line in log_messages:
-            f.write(line + "\n")
+class _MyLog():
+    def __init__(self):
+        self.log_messages = []
 
-def getlog():
-    return log_messages
+    def log(self, message):
+        self.log_messages.append(message)
 
-def clearlog():
-    log_messages.clear()
+    def getlog(self):
+        return self.log_messages
+
+    def clearlog(self):
+        self.log_messages.clear()
+
+    def setup_logging(self, level = logging.INFO):
+        root_logger = logging.getLogger()
+        root_logger.setLevel(level)
+
+        ch = logging.StreamHandler()
+        ch.setLevel(level)
+        ch.setFormatter(CustomFormatter())
+
+        list_handler = ListHandler(self.log_messages)
+        list_handler.setLevel(level)
+        list_handler.setFormatter(CustomFormatter())
+
+        root_logger.addHandler(ch)
+        root_logger.addHandler(list_handler)
+
+    def writelog(self):
+        # write log to file
+        with open(f"{logs_dir}/supermega.log", "w") as f:
+            for line in self.log_messages:
+                f.write(line + "\n")
 
 
-def setup_logging(level = logging.INFO):
-    root_logger = logging.getLogger()
-    root_logger.setLevel(level)
-
-    ch = logging.StreamHandler()
-    ch.setLevel(level)
-    ch.setFormatter(CustomFormatter())
-
-    list_handler = ListHandler(log_messages)
-    list_handler.setLevel(level)
-    list_handler.setFormatter(CustomFormatter())
-
-    root_logger.addHandler(ch)
-    root_logger.addHandler(list_handler)
-
-def clear_log():
-    log_messages.clear()
+MyLog: _MyLog = _MyLog()

supermega.py

@@ -17,7 +17,7 @@ from sender import scannerDetectsBytes
 from model.project import Project
 from model.settings import Settings
 from model.defs import *
-from log import setup_logging, writelog, clearlog
+from log import MyLog
 from utils import delete_all_files_in_directory
 
 def main():
@@ -109,7 +109,7 @@ def start(settings: Settings):
         delete_all_files_in_directory(f"{logs_dir}/")
     # And logs
     observer.reset()
-    clearlog()
+    MyLog.clearlog()
     exit_code = 0  # 0 = success
 
     # Load our input
@@ -146,7 +146,7 @@ def start(settings: Settings):
                 short_call_patching = project.settings.short_call_patching)
         except Exception as e:
             logger.error(f'Error compiling: {e}')
-            writelog()
+            MyLog.writelog()
             return 1
 
     # Assemble: Assemble .asm to .shc (ASM -> SHC)
@@ -158,7 +158,7 @@ def start(settings: Settings):
                 shellcode_out = main_shc_file)
         except Exception as e:
             logger.error("Error: Assembling failed: {}".format(e))
-            writelog()
+            MyLog.writelog()
             return 2
     #shutil.copy(main_shc_file, "working/build/shellcode.bin")
     
@@ -184,11 +184,11 @@ def start(settings: Settings):
         phases.injector.inject_exe(main_shc_file, settings, project)
     except PermissionError as e:
         logger.error(f'Error writing file: {e}')
-        writelog()
+        MyLog.writelog()
         return 2
     except Exception as e:
         logger.error(f'Error injecting: {e}')
-        writelog()
+        MyLog.writelog()
         return 3
     
     observer.add_code("exe_final", extract_code_from_exe_file_ep(settings.inject_exe_out, 300))
@@ -202,7 +202,7 @@ def start(settings: Settings):
                 scannerDetectsBytes(data, filename, useBrotli=True, verify=settings.verify)
             except Exception as e:
                 logger.error(f'Error scanning: {e}')
-                writelog()
+                MyLog.writelog()
                 return 4
     else:
         # Start/verify it at the end
@@ -219,7 +219,7 @@ def start(settings: Settings):
     if settings.cleanup_files_on_exit:
         clean_files()
 
-    writelog()
+    MyLog.writelog()
     return exit_code
 
 
@@ -265,5 +265,5 @@ def verify_shellcode(shc_name):
 
 
 if __name__ == "__main__":
-    setup_logging()
+    MyLog.setup_logging()
     main()

tester.py

@@ -5,7 +5,7 @@ from config import config
 from model.defs import *
 
 from model.settings import Settings
-from log import setup_logging
+from log import MyLog
 from supermega import start
 
 
@@ -56,5 +56,5 @@ def main():
 
 
 if __name__ == "__main__":
-    setup_logging(level=logging.WARN)
+    MyLog.setup_logging(level=logging.WARN)
     main()

web.py

@@ -6,11 +6,11 @@ from flask import Flask
 import logging
 
 from app.views import views
-from log import setup_logging, writelog
+from log import MyLog
 
 if __name__ == "__main__":
 	logging.getLogger('werkzeug').setLevel(logging.ERROR)
-	setup_logging()
+	MyLog.setup_logging()
 	parser = argparse.ArgumentParser()
 	parser.add_argument('--listenip', type=str, help='IP to listen on', default="0.0.0.0")
 	parser.add_argument('--listenport', type=int, help='Port to listen on', default=5001)