From 23992328d00a979132418102dcd3e3cf88be3fc6 Mon Sep 17 00:00:00 2001 From: Dobin Date: Thu, 28 Mar 2024 19:57:06 +0000 Subject: [PATCH] refactor: make log a class MyLog --- log.py | 61 +++++++++++++++++++++++++++++----------------------- supermega.py | 18 ++++++++-------- tester.py | 4 ++-- web.py | 4 ++-- 4 files changed, 47 insertions(+), 40 deletions(-) diff --git a/log.py b/log.py index d82b23e..a45ded4 100644 --- a/log.py +++ b/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() diff --git a/supermega.py b/supermega.py index 4f2dce1..3390bae 100644 --- a/supermega.py +++ b/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() diff --git a/tester.py b/tester.py index 5f0a205..a64aace 100644 --- a/tester.py +++ b/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() diff --git a/web.py b/web.py index 441abc1..cd17465 100644 --- a/web.py +++ b/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)