Merge pull request #229 from Pennyw0rth/marshall-logging-fix

Logging fixes (double logging & function caller obfuscation)

nxc/logger.py

@@ -9,6 +9,49 @@ from termcolor import colored
 from datetime import datetime
 from rich.text import Text
 from rich.logging import RichHandler
+import functools
+import inspect
+
+
+def create_temp_logger(caller_frame, formatted_text, args, kwargs):
+    """Create a temporary logger for emitting a log where we need to override the calling file & line number, since these are obfuscated"""
+    temp_logger = logging.getLogger("temp")
+    formatter = logging.Formatter("%(message)s", datefmt="[%X]")
+    handler = SmartDebugRichHandler(formatter=formatter)
+    handler.handle(LogRecord(temp_logger.name, logging.DEBUG, caller_frame.f_code.co_filename, caller_frame.f_lineno, formatted_text, args, kwargs, caller_frame=caller_frame))
+
+
+class SmartDebugRichHandler(RichHandler):
+    """Custom logging handler for when we want to log normal messages to DEBUG and not double log"""
+    def __init__(self, formatter=None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if formatter is not None:
+            self.setFormatter(formatter)
+            
+    def emit(self, record):
+        """Overrides the emit method of the RichHandler class so we can set the proper pathname and lineno"""
+        if hasattr(record, "caller_frame"):
+            frame_info = inspect.getframeinfo(record.caller_frame)
+            record.pathname = frame_info.filename
+            record.lineno = frame_info.lineno
+        super().emit(record)
+
+
+def no_debug(func):
+    """Stops logging non-debug messages when we are in debug mode
+    It creates a temporary logger and logs the message to the console and file
+    This is so we don't get both normal output AND debugging output, AND so we get the proper log calling file & line number
+    """
+    @functools.wraps(func)
+    def wrapper(self, msg, *args, **kwargs):
+        if self.logger.getEffectiveLevel() >= logging.INFO:
+            return func(self, msg, *args, **kwargs)
+        else:            
+            formatted_text = Text.from_ansi(self.format(msg, *args, **kwargs)[0])
+            caller_frame = inspect.currentframe().f_back
+            create_temp_logger(caller_frame, formatted_text, args, kwargs)
+            self.log_console_to_file(formatted_text, *args, **kwargs)
+    return wrapper
 
 
 class NXCAdapter(logging.LoggerAdapter):
@@ -54,6 +97,7 @@ class NXCAdapter(logging.LoggerAdapter):
 
         return (f"{module_name:<24} {self.extra['host']:<15} {self.extra['port']:<6} {self.extra['hostname'] if self.extra['hostname'] else 'NONE':<16} {msg}", kwargs)
 
+    @no_debug
     def display(self, msg, *args, **kwargs):
         """Display text to console, formatted for nxc"""
         msg, kwargs = self.format(f"{colored('[*]', 'blue', attrs=['bold'])} {msg}", kwargs)
@@ -61,13 +105,15 @@ class NXCAdapter(logging.LoggerAdapter):
         nxc_console.print(text, *args, **kwargs)
         self.log_console_to_file(text, *args, **kwargs)
 
+    @no_debug
     def success(self, msg, color="green", *args, **kwargs):
-        """Print some sort of success to the user"""
+        """Prints some sort of success to the user"""
         msg, kwargs = self.format(f"{colored('[+]', color, attrs=['bold'])} {msg}", kwargs)
         text = Text.from_ansi(msg)
         nxc_console.print(text, *args, **kwargs)
         self.log_console_to_file(text, *args, **kwargs)
 
+    @no_debug
     def highlight(self, msg, *args, **kwargs):
         """Prints a completely yellow highlighted message to the user"""
         msg, kwargs = self.format(f"{colored(msg, 'yellow', attrs=['bold'])}", kwargs)
@@ -75,6 +121,7 @@ class NXCAdapter(logging.LoggerAdapter):
         nxc_console.print(text, *args, **kwargs)
         self.log_console_to_file(text, *args, **kwargs)
 
+    @no_debug
     def fail(self, msg, color="red", *args, **kwargs):
         """Prints a failure (may or may not be an error) - e.g. login creds didn't work"""
         msg, kwargs = self.format(f"{colored('[-]', color, attrs=['bold'])} {msg}", kwargs)
@@ -88,16 +135,12 @@ class NXCAdapter(logging.LoggerAdapter):
         If debug or info logging is not enabled, we still want display/success/fail logged to the file specified,
         so we create a custom LogRecord and pass it to all the additional handlers (which will be all the file handlers)
         """
-        if self.logger.getEffectiveLevel() >= logging.INFO:
-            # will be 0 if it's just the console output, so only do this if we actually have file loggers
-            if len(self.logger.handlers):
-                try:
-                    for handler in self.logger.handlers:
-                        handler.handle(LogRecord("nxc", 20, "", kwargs, msg=text, args=args, exc_info=None))
-                except Exception as e:
-                    self.logger.fail(f"Issue while trying to custom print handler: {e}")
-        else:
-            self.logger.info(text)
+        if self.logger.getEffectiveLevel() >= logging.INFO and len(self.logger.handlers):  # will be 0 if it's just the console output, so only do this if we actually have file loggers
+            try:
+                for handler in self.logger.handlers:
+                    handler.handle(LogRecord("nxc", 20, "", kwargs, msg=text, args=args, exc_info=None))
+            except Exception as e:
+                self.logger.fail(f"Issue while trying to custom print handler: {e}")
 
     def add_file_log(self, log_file=None):
         file_formatter = TermEscapeCodeFormatter("%(asctime)s - %(levelname)s - %(message)s")
@@ -131,7 +174,7 @@ class NXCAdapter(logging.LoggerAdapter):
             datetime.now().strftime("%Y-%m-%d"),
             f"log_{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.log",
         )
-
+    
 
 class TermEscapeCodeFormatter(logging.Formatter):
     """A class to strip the escape codes for logging to files"""