possible improvement for read until handling

scrapli/channel/async_channel.py

@@ -2,7 +2,7 @@
 import asyncio
 import re
 import time
-from io import SEEK_END, BytesIO
+from io import BytesIO
 
 try:
     from contextlib import asynccontextmanager
@@ -143,8 +143,7 @@ async def _read_until_prompt(self, buf: bytes = b"") -> bytes:
             b = await self.read()
             read_buf.write(b)
 
-            read_buf.seek(-self._base_channel_args.comms_prompt_search_depth, SEEK_END)
-            search_buf = read_buf.read()
+            search_buf = self._process_read_buf(read_buf=read_buf)
 
             channel_match = re.search(
                 pattern=search_pattern,
@@ -185,8 +184,7 @@ async def _read_until_explicit_prompt(self, prompts: List[str]) -> bytes:
             b = await self.read()
             read_buf.write(b)
 
-            read_buf.seek(-self._base_channel_args.comms_prompt_search_depth, SEEK_END)
-            search_buf = read_buf.read()
+            search_buf = self._process_read_buf(read_buf=read_buf)
 
             for search_pattern in search_patterns:
                 channel_match = re.search(
@@ -247,8 +245,7 @@ async def _read_until_prompt_or_time(
             except ScrapliTimeout:
                 pass
 
-            read_buf.seek(-self._base_channel_args.comms_prompt_search_depth, SEEK_END)
-            search_buf = read_buf.read()
+            search_buf = self._process_read_buf(read_buf=read_buf)
 
             if (time.time() - start) > read_duration:
                 break

scrapli/channel/base_channel.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from functools import lru_cache
-from io import BytesIO
+from io import SEEK_END, BytesIO
 from typing import BinaryIO, List, Optional, Pattern, Tuple, Union
 
 from scrapli.exceptions import ScrapliAuthenticationFailed, ScrapliTypeError, ScrapliValueError
@@ -286,6 +286,35 @@ def close(self) -> None:
         if self.channel_log:
             self.channel_log.close()
 
+    def _process_read_buf(self, read_buf: BytesIO) -> bytes:
+        """
+        Process the read buffer
+
+        Seeks backwards up to search depth then partitions on newlines. Partition is to ensure that
+        the resulting search_buf does not end up with partial lines in the output which can cause
+        prompt patterns to match places they should not match!
+
+        Args:
+            read_buf: bytesio object read from the transport
+
+        Returns:
+            bytes: cleaned up search buffer
+
+        Raises:
+            N/A
+
+        """
+        read_buf.seek(-self._base_channel_args.comms_prompt_search_depth, SEEK_END)
+        search_buf = read_buf.read()
+
+        before, _, search_buf = search_buf.partition(b"\n")
+
+        if not search_buf:
+            # didn't split on anything or nothing after partition
+            search_buf = before
+
+        return search_buf
+
     def write(self, channel_input: str, redacted: bool = False) -> None:
         """
         Write input to the underlying Transport session

scrapli/channel/sync_channel.py

@@ -3,7 +3,7 @@
 import time
 from contextlib import contextmanager
 from datetime import datetime
-from io import SEEK_END, BytesIO
+from io import BytesIO
 from threading import Lock
 from typing import Iterator, List, Optional, Tuple
 
@@ -136,8 +136,7 @@ def _read_until_prompt(self, buf: bytes = b"") -> bytes:
         while True:
             read_buf.write(self.read())
 
-            read_buf.seek(-self._base_channel_args.comms_prompt_search_depth, SEEK_END)
-            search_buf = read_buf.read()
+            search_buf = self._process_read_buf(read_buf=read_buf)
 
             channel_match = re.search(
                 pattern=search_pattern,
@@ -177,8 +176,7 @@ def _read_until_explicit_prompt(self, prompts: List[str]) -> bytes:
         while True:
             read_buf.write(self.read())
 
-            read_buf.seek(-self._base_channel_args.comms_prompt_search_depth, SEEK_END)
-            search_buf = read_buf.read()
+            search_buf = self._process_read_buf(read_buf=read_buf)
 
             for search_pattern in search_patterns:
                 channel_match = re.search(
@@ -238,8 +236,7 @@ def _read_until_prompt_or_time(
             except ScrapliTimeout:
                 pass
 
-            read_buf.seek(-self._base_channel_args.comms_prompt_search_depth, SEEK_END)
-            search_buf = read_buf.read()
+            search_buf = self._process_read_buf(read_buf=read_buf)
 
             if (time.time() - start) > read_duration:
                 break

tests/unit/channel/test_async_channel.py

@@ -458,8 +458,8 @@ async def test_send_inputs_interact(monkeypatch, async_channel):
     _read_counter = 0
     _event_counter = 0
 
-    interact_events = [("clear logg", "[confirm]\n"), ("", "scrapli>")]
-    expected_buf = b"clear logg[confirm]\nscrapli>"
+    interact_events = [("clear logg", "[confirm]"), ("", "scrapli>")]
+    expected_buf = b"clear logg[confirm]scrapli>"
 
     async def _read(cls):
         nonlocal _read_counter

tests/unit/channel/test_base_channel.py

@@ -93,6 +93,35 @@ def test_channel_log_user_bytesio(base_transport_no_abc):
     assert chan.channel_log is bytes_log
 
 
+@pytest.mark.parametrize(
+    "test_data",
+    (
+        (
+            b"blah basic stuff>",
+            b"blah basic stuff>",
+        ),
+        (
+            b"\noper-status></physical-interface>\nsomeprompt>",
+            b"oper-status></physical-interface>\nsomeprompt>",
+        ),
+    ),
+    ids=("simple_buf", "xml_out"),
+)
+def test_process_read_buf(test_data, base_channel):
+    """
+    This asserts that the process read buf method always returns a search buf that is "roooted"
+    on newlines -- meaning we never scan backwards through the readbuf and return a line that gets
+    split up resulting in a section of the line being at the "start" of the search_buf that looks
+    exactly like a normal prompt we would match on
+    """
+    inbuf, expected_buf = test_data
+
+    read_buf = BytesIO(inbuf)
+    search_buf = base_channel._process_read_buf(read_buf=read_buf)
+
+    assert search_buf == expected_buf
+
+
 def test_channel_write(caplog, monkeypatch, base_channel):
     caplog.set_level(logging.DEBUG, logger="scrapli.channel")
 

tests/unit/channel/test_sync_channel.py

@@ -422,8 +422,8 @@ def test_send_inputs_interact(monkeypatch, sync_channel):
     _read_counter = 0
     _event_counter = 0
 
-    interact_events = [("clear logg", "[confirm]\n"), ("", "scrapli>")]
-    expected_buf = b"clear logg[confirm]\nscrapli>"
+    interact_events = [("clear logg", "[confirm]"), ("", "scrapli>")]
+    expected_buf = b"clear logg[confirm]scrapli>"
 
     def _read(cls):
         nonlocal _read_counter