From 2524547ccfe071510dd7099153be3055f2d3d4ef Mon Sep 17 00:00:00 2001
From: Dobin <dobin.rutishauser@gmail.com>
Date: Sat, 17 Feb 2024 15:01:38 +0000
Subject: [PATCH] feature: add logging output too to logs/

---
 r2helper.py  |  2 +-
 supermega.py | 99 +++++++++++++++++++++++++++++++++-------------------
 2 files changed, 65 insertions(+), 36 deletions(-)

diff --git a/r2helper.py b/r2helper.py
index f1b626e..03ae184 100644
--- a/r2helper.py
+++ b/r2helper.py
@@ -19,7 +19,7 @@ def r2_disas(data: bytes):
         f.write(data)
     code_len = len(data)
 
-    r2 = r2pipe.open(filename)
+    r2 = r2pipe.open(filename, flags=['-e', 'scr.prompt=false', '-e'])
     r2.cmd('aaa')
 
     r2.cmd('e scr.color=0')
diff --git a/supermega.py b/supermega.py
index 4647aca..9f88158 100644
--- a/supermega.py
+++ b/supermega.py
@@ -18,40 +18,7 @@ from observer import observer
 from project import project
 
 
-# ANSI escape sequences for colors
-class LogColors:
-    HEADER = '\033[95m'
-    BLUE = '\033[94m'
-    GREEN = '\033[92m'
-    WARNING = '\033[93m'
-    FAIL = '\033[91m'
-    ENDC = '\033[0m'
-    BOLD = '\033[1m'
-    UNDERLINE = '\033[4m'
-
-# Custom formatter to include colors in log output
-class CustomFormatter(logging.Formatter):
-    #format = "%(asctime)s - %(name)-12s - [%(levelname)-8s] - %(message)s (%(filename)s:%(lineno)d)"
-    format = "(%(filename)-12s) %(message)s"
-
-    FORMATS = {
-        logging.DEBUG: format,
-        logging.INFO: format,
-        logging.WARNING: LogColors.WARNING + format + LogColors.ENDC,
-        logging.ERROR: LogColors.FAIL + format + LogColors.ENDC,
-        logging.CRITICAL: LogColors.FAIL + LogColors.BOLD + format + LogColors.ENDC
-    }
-
-    def format(self, record):
-        log_fmt = self.FORMATS.get(record.levelno)
-        formatter = logging.Formatter(log_fmt, datefmt="%Y-%m-%d %H:%M:%S")
-        return formatter.format(record)
-
-# Configure logging
-handler = logging.StreamHandler()
-handler.setFormatter(CustomFormatter())
-logging.basicConfig(level=logging.DEBUG, handlers=[handler])
-logger = logging.getLogger("ExploitLogger")
+log_messages = []
 
 
 def main():
@@ -255,6 +222,11 @@ def start():
     if project.cleanup_files_on_exit:
         clean_files()
 
+    # write log to file
+    with open("logs/supermega.log", "w") as f:
+        for line in log_messages:
+            f.write(line + "\n")
+
     exit(exit_code)
 
 
@@ -307,7 +279,64 @@ def verify_shellcode(shc_name):
     else:
         logger.info("---> Verify FAIL. Shellcode doesnt work (file was not created)")
         return False
-    
+
+
+# Logging
+
+# ANSI escape sequences for colors
+class LogColors:
+    HEADER = '\033[95m'
+    BLUE = '\033[94m'
+    GREEN = '\033[92m'
+    WARNING = '\033[93m'
+    FAIL = '\033[91m'
+    ENDC = '\033[0m'
+    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"
+
+    FORMATS = {
+        logging.DEBUG: format,
+        logging.INFO: format,
+        logging.WARNING: LogColors.WARNING + format + LogColors.ENDC,
+        logging.ERROR: LogColors.FAIL + format + LogColors.ENDC,
+        logging.CRITICAL: LogColors.FAIL + LogColors.BOLD + format + LogColors.ENDC
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        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__()
+        self.log_list = log_list
+
+    def emit(self, record):
+        # Format the log record and store it in the list
+        log_entry = self.format(record)
+        self.log_list.append(log_entry)
+
+def setup_logging():
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.DEBUG)
+
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.DEBUG)
+    ch.setFormatter(CustomFormatter())
+
+    list_handler = ListHandler(log_messages)
+    list_handler.setLevel(logging.DEBUG)
+    list_handler.setFormatter(CustomFormatter())
+
+    root_logger.addHandler(ch)
+    root_logger.addHandler(list_handler)
+
 
 if __name__ == "__main__":
+    setup_logging()
     main()