@@ -12,6 +12,7 @@ def __init__(
12
12
file_encoding : str = "utf-8" ,
13
13
no_log : Optional [Dict [str , Any ]] = None ,
14
14
record_writes : bool = False ,
15
+ slog_buffer : Optional [io .StringIO ] = None ,
15
16
) -> None :
16
17
if no_log is None :
17
18
self .no_log = {}
@@ -30,6 +31,13 @@ def __init__(
30
31
else :
31
32
self .session_log = None
32
33
34
+ # In order to ensure all the no_log entries get hidden properly
35
+ # we must first store everying in memory and then write out to file.
36
+ # Otherwise, we might miss the data we are supposed to hide (they span
37
+ # multiple reads).
38
+ if slog_buffer is None :
39
+ self .slog_buffer = io .StringIO ()
40
+
33
41
# Ensures last write operations prior to disconnect are recorded.
34
42
self .fin = False
35
43
@@ -49,15 +57,24 @@ def open(self) -> None:
49
57
50
58
def close (self ) -> None :
51
59
"""Close the session_log file (if it is a file that we opened)."""
60
+ self .flush ()
52
61
if self .session_log and self ._session_log_close :
53
62
self .session_log .close ()
54
63
self .session_log = None
55
64
56
- def write (self , data : str ) -> None :
57
- if self .session_log is not None and len (data ) > 0 :
58
- # Hide the password and secret in the session_log
59
- for hidden_data in self .no_log .values ():
60
- data = data .replace (hidden_data , "********" )
65
+ def no_log_filter (self , data : str ) -> str :
66
+ """Filter content from the session_log."""
67
+ for hidden_data in self .no_log .values ():
68
+ data = data .replace (hidden_data , "********" )
69
+ return data
70
+
71
+ def flush (self ) -> None :
72
+ """Force the slog_buffer to be written out to the actual file"""
73
+
74
+ if self .session_log is not None :
75
+ self .slog_buffer .seek (0 )
76
+ data = self .slog_buffer .read ()
77
+ data = self .no_log_filter (data )
61
78
62
79
if isinstance (self .session_log , io .BufferedIOBase ):
63
80
self .session_log .write (write_bytes (data , encoding = self .file_encoding ))
@@ -67,4 +84,10 @@ def write(self, data: str) -> None:
67
84
assert isinstance (self .session_log , io .BufferedIOBase ) or isinstance (
68
85
self .session_log , io .TextIOBase
69
86
)
87
+
88
+ # Flush the underlying file
70
89
self .session_log .flush ()
90
+
91
+ def write (self , data : str ) -> None :
92
+ if len (data ) > 0 :
93
+ self .slog_buffer .write (data )
0 commit comments