From 7dec0010fdc251bf44ccaebe8e885816e0282595 Mon Sep 17 00:00:00 2001
From: Kirk Byers <ktbyers@twb-tech.com>
Date: Mon, 30 Oct 2023 11:32:36 -0700
Subject: [PATCH] First pass on session_log fixes

---
 netmiko/session_log.py            | 33 ++++++++++++++++++++++++++-----
 tests/test_netmiko_session_log.py |  5 +++++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/netmiko/session_log.py b/netmiko/session_log.py
index 9ead1316c..4a2cba10c 100644
--- a/netmiko/session_log.py
+++ b/netmiko/session_log.py
@@ -12,6 +12,7 @@ def __init__(
         file_encoding: str = "utf-8",
         no_log: Optional[Dict[str, Any]] = None,
         record_writes: bool = False,
+        slog_buffer: Optional[io.StringIO] = None,
     ) -> None:
         if no_log is None:
             self.no_log = {}
@@ -30,6 +31,13 @@ def __init__(
         else:
             self.session_log = None
 
+        # In order to ensure all the no_log entries get hidden properly
+        # we must first store everying in memory and then write out to file.
+        # Otherwise, we might miss the data we are supposed to hide (they span
+        # multiple reads).
+        if slog_buffer is None:
+            self.slog_buffer = io.StringIO()
+
         # Ensures last write operations prior to disconnect are recorded.
         self.fin = False
 
@@ -49,15 +57,24 @@ def open(self) -> None:
 
     def close(self) -> None:
         """Close the session_log file (if it is a file that we opened)."""
+        self.flush()
         if self.session_log and self._session_log_close:
             self.session_log.close()
             self.session_log = None
 
-    def write(self, data: str) -> None:
-        if self.session_log is not None and len(data) > 0:
-            # Hide the password and secret in the session_log
-            for hidden_data in self.no_log.values():
-                data = data.replace(hidden_data, "********")
+    def no_log_filter(self, data: str) -> str:
+        """Filter content from the session_log."""
+        for hidden_data in self.no_log.values():
+            data = data.replace(hidden_data, "********")
+        return data
+
+    def flush(self) -> None:
+        """Force the slog_buffer to be written out to the actual file"""
+
+        if self.session_log is not None:
+            self.slog_buffer.seek(0)
+            data = self.slog_buffer.read()
+            data = self.no_log_filter(data)
 
             if isinstance(self.session_log, io.BufferedIOBase):
                 self.session_log.write(write_bytes(data, encoding=self.file_encoding))
@@ -67,4 +84,10 @@ def write(self, data: str) -> None:
             assert isinstance(self.session_log, io.BufferedIOBase) or isinstance(
                 self.session_log, io.TextIOBase
             )
+
+            # Flush the underlying file
             self.session_log.flush()
+
+    def write(self, data: str) -> None:
+        if len(data) > 0:
+            self.slog_buffer.write(data)
diff --git a/tests/test_netmiko_session_log.py b/tests/test_netmiko_session_log.py
index bc41398ab..ef8469524 100644
--- a/tests/test_netmiko_session_log.py
+++ b/tests/test_netmiko_session_log.py
@@ -171,6 +171,8 @@ def test_unicode(device_slog):
     conn.session_log.write("\nTesting unicode\n")
     conn.session_log.write(smiley_face)
     conn.session_log.write(smiley_face)
+    time.sleep(1)
+    conn.session_log.flush()
 
     file_name = device_slog["session_log"]
     with open(file_name, "r") as f:
@@ -241,3 +243,6 @@ def test_session_log_custom_secrets(device_slog):
         assert sanitize_secrets["secret3"] not in session_log
     if sanitize_secrets.get("supersecret") is not None:
         assert sanitize_secrets["supersecret"] not in session_log
+
+
+def test_session_log_no_log(device_slog):