add buffering to Tube

ptrlib/pwn/proc.py

@@ -64,7 +64,7 @@ def _settimeout(self, timeout):
 
     def _socket(self):
         return self.proc
-    
+
     def _poll(self):
         if self.proc is None:
             return False
@@ -126,7 +126,7 @@ def recv(self, size=4096, timeout=None):
             return None
 
         self._poll() # poll after received all data
-        
+
         if len(self.reservoir) == 0:
             # No data received
             data = None
@@ -140,41 +140,6 @@ def recv(self, size=4096, timeout=None):
             self.reservoir = b''
         return data
 
-    def recvonce(self, size=4, timeout=None):
-        """Receive raw data
-
-        Receive raw data of `size` bytes length through the pipe.
-
-        Args:
-            size    (int): The data size to receive
-            timeout (int): Timeout (in second)
-
-        Returns:
-            bytes: The received data
-
-        Raises:
-            SocketException: If the socket is broken.
-        """
-        self._settimeout(timeout)
-        data = b''
-        if size <= 0:
-            logger.error("`size` must be larger than 0")
-            return None
-
-        read_byte = 0
-        recv_size = size
-        while read_byte < size:
-            recv_data = self.recv(recv_size, timeout)
-            if recv_data is None:
-                return None
-            elif recv_data == b'':
-                logger.error("Received nothing")
-                return None
-            data += recv_data
-            read_byte += len(data)
-            recv_size = size - read_byte
-        return data
-
     def send(self, data, timeout=None):
         """Send raw data
 
@@ -215,7 +180,7 @@ def shutdown(self, target):
         """
         if target in ['write', 'send', 'stdin']:
             self.proc.stdin.close()
-        
+
         elif target in ['read', 'recv', 'stdout', 'stderr']:
             self.proc.stdout.close()
 

ptrlib/pwn/sock.py

@@ -40,8 +40,8 @@ def _settimeout(self, timeout):
 
     def _socket(self):
         return self.sock
-    
-    def recv(self, size=4096, timeout=None):
+
+    def _recv(self, size=4096, timeout=None):
         """Receive raw data
 
         Receive raw data of maximum `size` bytes length through the socket.
@@ -66,35 +66,6 @@ def recv(self, size=4096, timeout=None):
             data = None
         return data
 
-    def recvonce(self, size=4, timeout=None):
-        """Receive raw data at once
-
-        Receive raw data of `size` bytes length through the socket.
-
-        Args:
-            size    (int): The data size to receive
-            timeout (int): Timeout (in second)
-
-        Returns:
-            bytes: The received data
-        """
-        self._settimeout(timeout)
-        data = b''
-        if size <= 0:
-            logger.error("`size` must be larger than 0")
-            return None
-        try:
-            read_byte = 0
-            recv_size = size
-            while read_byte < size:
-                data += self.sock.recv(recv_size)
-                read_byte = len(data)
-                recv_size = size - read_byte
-        except socket.timeout:
-            logger.error("Timeout")
-            return None
-        return data
-
     def send(self, data, timeout=None):
         """Send raw data
 
@@ -134,7 +105,7 @@ def shutdown(self, target):
         """
         if target in ['write', 'send', 'stdin']:
             self.sock.shutdown(socket.SHUT_WR)
-        
+
         elif target in ['read', 'recv', 'stdout', 'stderr']:
             self.sock.shutdown(socket.SHUT_RD)
 

ptrlib/pwn/tube.py

@@ -9,105 +9,115 @@
 logger = getLogger(__name__)
 
 class Tube(metaclass=ABCMeta):
+    def __init__(self):
+        self.buf = b''
+
     @abstractmethod
     def _settimeout(self, timeout):
         pass
 
     @abstractmethod
-    def recv(self, size, timeout):
-        pass
+    def _recv(self, size, timeout):
+        """Receive raw data
 
-    @abstractmethod
-    def recvonce(self, size, timeout):
+        Receive raw data of maximum `size` bytes length through the socket.
+
+        Args:
+            size    (int): The data size to receive
+            timeout (int): Timeout (in second)
+
+        Returns:
+            bytes: The received data
+        """
         pass
 
-    def recvall(self, size=4096, timeout=None):
-        """Receive all data
+    def unget(self, data):
+        self.buf = data + self.buf
 
-        Receive all data through the socket.
+    def recv(self, size, timeout):
+        """Receive raw data with buffering
+
+        Receive raw data of maximum `size` bytes length through the socket.
 
         Args:
-            size (int)   : Data size to receive at once
+            size    (int): The data size to receive
             timeout (int): Timeout (in second)
 
         Returns:
             bytes: The received data
         """
-        data = b''
-        while True:
-            part = self.recv(size)
-            data += part
+        if size > len(self.bufsize):
+            self.buf += self._recv(size, timeout)
+
+        data, self.buf = self.buf[:size], self.buf[size:]
         return data
 
-    def recvline(self, timeout=None, drop=True):
-        """Receive a line
+    def recvonce(self, size, timeout):
+        """Receive raw data with buffering
 
-        Receive a line of raw data through the socket.
+        Receive raw data of `size` bytes length through the socket.
 
         Args:
+            size    (int): The data size to receive
             timeout (int): Timeout (in second)
-            drop (bool)  : Whether or not to strip the newline
 
         Returns:
             bytes: The received data
         """
         data = b''
-        c = None
-        while c != b'\n':
-            c = self.recvonce(1, timeout)
-            if c is None:
-                # Timeout
-                break
-            else:
-                data += c
-        if drop:
-            return data.rstrip()
-        else:
-            return data
+        while len(data) < size:
+            data += self.recv(size - len(data))
 
-    def recvuntil(self, delim, timeout=None):
+        if len(data) > size:
+            self.unget(data[size:])
+        return data[:size]
+
+    def recvall(self, size=4096, timeout=None):
+        """Receive all data
+
+        Receive all data through the socket.
+
+        Args:
+            size (int)   : Data size to receive at once
+            timeout (int): Timeout (in second)
+
+        Returns:
+            bytes: The received data
+        """
+        data = b''
+        while True:
+            part = self.recv(size)
+            data += part
+        return data
+
+
+    def recvuntil(self, size=4096, delim, timeout=None):
         """Receive raw data until `delim` comes
 
         Args:
+            size (int)   : The data size to receive at once
             delim (bytes): The delimiter bytes
             timeout (int): Timeout (in second)
 
         Returns:
             bytes: The received data
         """
+
         if isinstance(delim, str):
             delim = str2bytes(delim)
         data = b''
-        length = len(delim)
 
-        # Create the Boyer-Moore table
-        bm_table = [length for i in range(0x100)]
-        for (i, c) in enumerate(delim):
-            bm_table[c] = length - i - 1
+        while data.find(delim) == -1:
+            data += self.recv(size, timeout)
+        pos = data.find(delim) + len(delim)
+        self.unget(data[pos:])
+        return data[:pos]
 
-        # Receive data until the delimiter comes
-        recv_size = length
-        obj = None
-        while True:
-            # Receive
-            obj = self.recvonce(recv_size, timeout)
-            if obj is None:
-                # Timeout
-                break
-            else:
-                data += obj
-            # Search
-            i = -1
-            j = length - 1
-            while j >= 0:
-                if data[i] != delim[j]: break
-                i, j = i - 1, j - 1
-            if j < 0:
-                # Delimiter found
-                break
-            recv_size = max(bm_table[data[i]], length - j)
-            i += recv_size
-        return data
+    def recvline(self, size=4096, timeout=None, drop=True):
+        line = self.recvuntil(b'\n')
+        if drop:
+            return line.rstrip()
+        return line
 
     @abstractmethod
     def send(self, data, timeout):