99from datetime import datetime
1010from rich .text import Text
1111from rich .logging import RichHandler
12+ import functools
13+ import inspect
14+
15+
16+ def create_temp_logger (caller_frame , formatted_text , args , kwargs ):
17+ """Create a temporary logger for emitting a log where we need to override the calling file & line number, since these are obfuscated"""
18+ temp_logger = logging .getLogger ("temp" )
19+ formatter = logging .Formatter ("%(message)s" , datefmt = "[%X]" )
20+ handler = SmartDebugRichHandler (formatter = formatter )
21+ handler .handle (LogRecord (temp_logger .name , logging .INFO , caller_frame .f_code .co_filename , caller_frame .f_lineno , formatted_text , args , kwargs , caller_frame = caller_frame ))
22+
23+
24+ class SmartDebugRichHandler (RichHandler ):
25+ """Custom logging handler for when we want to log normal messages to DEBUG and not double log"""
26+ def __init__ (self , formatter = None , * args , ** kwargs ):
27+ super ().__init__ (* args , ** kwargs )
28+ if formatter is not None :
29+ self .setFormatter (formatter )
30+
31+ def emit (self , record ):
32+ """Overrides the emit method of the RichHandler class so we can set the proper pathname and lineno"""
33+ if hasattr (record , "caller_frame" ):
34+ frame_info = inspect .getframeinfo (record .caller_frame )
35+ record .pathname = frame_info .filename
36+ record .lineno = frame_info .lineno
37+ super ().emit (record )
38+
39+
40+ def no_debug (func ):
41+ """Stops logging non-debug messages when we are in debug mode
42+ It creates a temporary logger and logs the message to the console and file
43+ This is so we don't get both normal output AND debugging output, AND so we get the proper log calling file & line number
44+ """
45+ @functools .wraps (func )
46+ def wrapper (self , msg , * args , ** kwargs ):
47+ if self .logger .getEffectiveLevel () >= logging .INFO :
48+ return func (self , msg , * args , ** kwargs )
49+ else :
50+ formatted_text = Text .from_ansi (self .format (msg , * args , ** kwargs )[0 ])
51+ caller_frame = inspect .currentframe ().f_back
52+ create_temp_logger (caller_frame , formatted_text , args , kwargs )
53+ self .log_console_to_file (formatted_text , * args , ** kwargs )
54+ return wrapper
1255
1356
1457class NXCAdapter (logging .LoggerAdapter ):
1558 def __init__ (self , extra = None ):
1659 logging .basicConfig (
1760 format = "%(message)s" ,
1861 datefmt = "[%X]" ,
19- handlers = [
20- RichHandler (
21- console = nxc_console ,
22- rich_tracebacks = True ,
23- tracebacks_show_locals = False ,
24- )
25- ],
62+ handlers = [RichHandler (
63+ console = nxc_console ,
64+ rich_tracebacks = True ,
65+ tracebacks_show_locals = False
66+ )],
2667 )
2768 self .logger = logging .getLogger ("nxc" )
2869 self .extra = extra
@@ -40,52 +81,47 @@ def format(self, msg, *args, **kwargs): # noqa: A003
4081 if self .extra is None :
4182 return f"{ msg } , kwargs
4283
43- if "module_name" in self .extra and len (self .extra ["module_name" ]) > 8 :
84+ if "module_name" in self .extra and len (self .extra ["module_name" ]) > 11 :
4485 self .extra ["module_name" ] = self .extra ["module_name" ][:8 ] + "..."
4586
4687 # If the logger is being called when hooking the 'options' module function
4788 if len (self .extra ) == 1 and ("module_name" in self .extra ):
48- return (
49- f"{ colored (self .extra ['module_name' ], 'cyan' , attrs = ['bold' ]):<64} { msg } ,
50- kwargs ,
51- )
89+ return (f"{ colored (self .extra ['module_name' ], 'cyan' , attrs = ['bold' ]):<64} { msg } , kwargs )
5290
5391 # If the logger is being called from nxcServer
5492 if len (self .extra ) == 2 and ("module_name" in self .extra ) and ("host" in self .extra ):
55- return (
56- f"{ colored (self .extra ['module_name' ], 'cyan' , attrs = ['bold' ]):<24} { self .extra ['host' ]:<39} { msg } ,
57- kwargs ,
58- )
93+ return (f"{ colored (self .extra ['module_name' ], 'cyan' , attrs = ['bold' ]):<24} { self .extra ['host' ]:<39} { msg } , kwargs )
5994
6095 # If the logger is being called from a protocol
6196 module_name = colored (self .extra ["module_name" ], "cyan" , attrs = ["bold" ]) if "module_name" in self .extra else colored (self .extra ["protocol" ], "blue" , attrs = ["bold" ])
6297
63- return (
64- f"{ module_name :<24} { self .extra ['host' ]:<15} { self .extra ['port' ]:<6} { self .extra ['hostname' ] if self .extra ['hostname' ] else 'NONE' :<16} { msg } ,
65- kwargs ,
66- )
98+ 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 )
6799
100+ @no_debug
68101 def display (self , msg , * args , ** kwargs ):
69102 """Display text to console, formatted for nxc"""
70103 msg , kwargs = self .format (f"{ colored ('[*]' , 'blue' , attrs = ['bold' ])} { msg } , kwargs )
71104 text = Text .from_ansi (msg )
72105 nxc_console .print (text , * args , ** kwargs )
73106 self .log_console_to_file (text , * args , ** kwargs )
74107
108+ @no_debug
75109 def success (self , msg , color = "green" , * args , ** kwargs ):
76- """Print some sort of success to the user"""
110+ """Prints some sort of success to the user"""
77111 msg , kwargs = self .format (f"{ colored ('[+]' , color , attrs = ['bold' ])} { msg } , kwargs )
78112 text = Text .from_ansi (msg )
79113 nxc_console .print (text , * args , ** kwargs )
80114 self .log_console_to_file (text , * args , ** kwargs )
81115
116+ @no_debug
82117 def highlight (self , msg , * args , ** kwargs ):
83118 """Prints a completely yellow highlighted message to the user"""
84119 msg , kwargs = self .format (f"{ colored (msg , 'yellow' , attrs = ['bold' ])} , kwargs )
85120 text = Text .from_ansi (msg )
86121 nxc_console .print (text , * args , ** kwargs )
87122 self .log_console_to_file (text , * args , ** kwargs )
88123
124+ @no_debug
89125 def fail (self , msg , color = "red" , * args , ** kwargs ):
90126 """Prints a failure (may or may not be an error) - e.g. login creds didn't work"""
91127 msg , kwargs = self .format (f"{ colored ('[-]' , color , attrs = ['bold' ])} { msg } , kwargs )
@@ -99,26 +135,12 @@ def log_console_to_file(self, text, *args, **kwargs):
99135 If debug or info logging is not enabled, we still want display/success/fail logged to the file specified,
100136 so we create a custom LogRecord and pass it to all the additional handlers (which will be all the file handlers)
101137 """
102- if self .logger .getEffectiveLevel () >= logging .INFO :
103- # will be 0 if it's just the console output, so only do this if we actually have file loggers
104- if len (self .logger .handlers ):
105- try :
106- for handler in self .logger .handlers :
107- handler .handle (
108- LogRecord (
109- "nxc" ,
110- 20 ,
111- "" ,
112- kwargs ,
113- msg = text ,
114- args = args ,
115- exc_info = None ,
116- )
117- )
118- except Exception as e :
119- self .logger .fail (f"Issue while trying to custom print handler: { e } )
120- else :
121- self .logger .info (text )
138+ 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
139+ try :
140+ for handler in self .logger .handlers :
141+ handler .handle (LogRecord ("nxc" , 20 , "" , kwargs , msg = text , args = args , exc_info = None ))
142+ except Exception as e :
143+ self .logger .fail (f"Issue while trying to custom print handler: { e } )
122144
123145 def add_file_log (self , log_file = None ):
124146 file_formatter = TermEscapeCodeFormatter ("%(asctime)s - %(levelname)s - %(message)s" )
@@ -152,7 +174,7 @@ def init_log_file():
152174 datetime .now ().strftime ("%Y-%m-%d" ),
153175 f"log_{ datetime .now ().strftime ('%Y-%m-%d-%H-%M-%S' )} ,
154176 )
155-
177+
156178
157179class TermEscapeCodeFormatter (logging .Formatter ):
158180 """A class to strip the escape codes for logging to files"""
0 commit comments