Variable Size Buffer Fill (#806)

Support custom buffer-read sizes for `tube` objects via their internal `Buffer` object.  Default buffer size is 4096.  Buffer size can be set per-object at instantiation-time (`tube(buffer_fill_size=...)`), by attribute (`tube.buffer_fill_size=...`), or in a context (`context.buffer_size=...`).

pwnlib/commandline/main.py

@@ -16,8 +16,8 @@
 from . import shellcraft
 from . import unhex
 from . import update
-from .common import parser
 from ..context import context
+from .common import parser
 
 commands = {
     'asm': asm.main,

pwnlib/context/__init__.py

@@ -329,6 +329,7 @@ class ContextType(object):
         'aslr': True,
         'binary': None,
         'bits': 32,
+        'buffer_size': 4096,
         'device': os.getenv('ANDROID_SERIAL', None) or None,
         'endian': 'little',
         'kernel': None,
@@ -1140,6 +1141,14 @@ def adb(self):
 
         return command
 
+    @_validator
+    def buffer_size(self, size):
+        """Internal buffer size to use for ``tube`` objects.
+
+        This is not the maximum size of the buffer, but this is the amount of data
+        which is passed to each raw ``read`` syscall (or equivalent).
+        """
+        return int(size)
 
     #*************************************************************************
     #                               ALIASES

pwnlib/data/useragents/download-useragents.py

@@ -7,6 +7,7 @@
 import urllib
 
 from bs4 import BeautifulSoup
+
 from pwn import *
 
 uas = set()

pwnlib/term/term.py

@@ -21,8 +21,8 @@
 import threading
 import traceback
 
-from ..context import ContextType
 from . import termcap
+from ..context import ContextType
 
 settings = None
 _graphics_mode = False

pwnlib/tubes/buffer.py

@@ -1,5 +1,8 @@
 #!/usr/bin/env python2
 
+from ..context import context
+
+
 class Buffer(Exception):
     """
     List of strings with some helper routines.
@@ -28,10 +31,10 @@ class Buffer(Exception):
         The ``0th`` item in the buffer is the oldest item, and
         will be received first.
     """
-    def __init__(self):
+    def __init__(self, buffer_fill_size = None):
         self.data = [] # Buffer
         self.size = 0  # Length
-
+        self.buffer_fill_size = buffer_fill_size
 
     def __len__(self):
         """
@@ -169,3 +172,19 @@ def get(self, want=float('inf')):
         self.size -= len(data)
 
         return data
+
+    def get_fill_size(self,size=None):
+        """
+        Retrieves the default fill size for this buffer class.
+
+        Arguments:
+            size (int): (Optional) If set and not None, returns the size variable back.
+
+        Returns:
+            Fill size as integer if size == None, else size.
+        """
+        if size:
+            return size
+
+        with context.local(buffer_size = self.buffer_fill_size):
+            return context.buffer_size

pwnlib/tubes/listen.py

@@ -20,13 +20,11 @@ class listen(sock):
         bindaddr(str): The address to bind to.
         fam: The string "any", "ipv4" or "ipv6" or an integer to pass to :func:`socket.getaddrinfo`.
         typ: The string "tcp" or "udp" or an integer to pass to :func:`socket.getaddrinfo`.
-        timeout: A positive number, None
     """
 
     def __init__(self, port=0, bindaddr = "0.0.0.0",
-                 fam = "any", typ = "tcp",
-                 timeout = Timeout.default, level = None):
-        super(listen, self).__init__(timeout, level = level)
+                 fam = "any", typ = "tcp", *args, **kwargs):
+        super(listen, self).__init__(*args, **kwargs)
 
         port = int(port)
         fam  = {socket.AF_INET: 'ipv4',

pwnlib/tubes/process.py

@@ -43,8 +43,6 @@ class process(tube):
             Working directory.  Uses the current working directory by default.
         env(dict):
             Environment variables.  By default, inherits from Python's environment.
-        timeout(int):
-            Timeout to use on ``tube`` ``recv`` operations.
         stdin(int):
             File object or file descriptor number to use for ``stdin``.
             By default, a pipe is used.  A pty can be used instead by setting
@@ -202,20 +200,21 @@ def __init__(self, argv = None,
                  executable = None,
                  cwd = None,
                  env = None,
-                 timeout = Timeout.default,
                  stdin  = PIPE,
                  stdout = PTY,
                  stderr = STDOUT,
-                 level = None,
                  close_fds = True,
                  preexec_fn = lambda: None,
                  raw = True,
                  aslr = None,
                  setuid = None,
                  where = 'local',
                  display = None,
-                 alarm = None):
-        super(process, self).__init__(timeout, level = level)
+                 alarm = None,
+                 *args,
+                 **kwargs
+                 ):
+        super(process, self).__init__(*args,**kwargs)
 
         # Permit using context.binary
         if argv is None:

pwnlib/tubes/remote.py

@@ -45,8 +45,8 @@ class remote(sock):
 
     def __init__(self, host, port,
                  fam = "any", typ = "tcp",
-                 timeout = Timeout.default, ssl=False, sock=None, level = None):
-        super(remote, self).__init__(timeout, level = level)
+                 ssl=False, sock=None, *args, **kwargs):
+        super(remote, self).__init__(*args, **kwargs)
 
         self.rport  = int(port)
         self.rhost  = host
@@ -158,17 +158,13 @@ def fromsocket(cls, socket):
 
 class tcp(remote):
     __doc__ = remote.__doc__
-    def __init__(self, host, port,
-                 fam = "any", typ = "tcp",
-                 timeout = Timeout.default, ssl=False, sock=None, level = None):
-        return super(tcp, self).__init__(host, port, fam, typ, timeout, ssl, sock, level)
+    def __init__(self, host, port, *a, **kw):
+        return super(tcp, self).__init__(host, port, typ="tcp", *a, **kw)
 
 class udp(remote):
     __doc__ = remote.__doc__
-    def __init__(self, host, port,
-                 fam = "any", typ = "udp",
-                 timeout = Timeout.default, ssl=False, sock=None, level = None):
-        return super(udp, self).__init__(host, port, fam, typ, timeout, ssl, sock, level)
+    def __init__(self, host, port, *a, **kw):
+        return super(udp, self).__init__(host, port, typ="udp", *a, **kw)
 
 class connect(remote):
     __doc__ = remote.__doc__

pwnlib/tubes/serialtube.py

@@ -18,10 +18,8 @@ def __init__(
             self, port = None, baudrate = 115200,
             convert_newlines = True,
             bytesize = 8, parity='N', stopbits=1, xonxoff = False,
-            rtscts = False, dsrdtr = False,
-            timeout = Timeout.default,
-            level = None):
-        super(serialtube, self).__init__(timeout, level = level)
+            rtscts = False, dsrdtr = False):
+        super(serialtube, self).__init__(*a, **kw)
 
         if port is None:
             if platform.system() == 'Darwin':

pwnlib/tubes/sock.py

@@ -10,8 +10,8 @@
 class sock(tube):
     """Methods available exclusively to sockets."""
 
-    def __init__(self, timeout, level = None):
-        super(sock, self).__init__(timeout, level = level)
+    def __init__(self, *args, **kwargs):
+        super(sock, self).__init__(*args, **kwargs)
         self.closed = {"recv": False, "send": False}
 
     # Overwritten for better usability

pwnlib/tubes/ssh.py

@@ -59,16 +59,16 @@ class ssh_channel(sock):
     #: Only valid when instantiated through :meth:`ssh.process`
     pid = None
 
-    #: Executable of the process
+    #: Executable of the procesks
     #: Only valid when instantiated through :meth:`ssh.process`
     executable = None
 
     #: Arguments passed to the process
     #: Only valid when instantiated through :meth:`ssh.process`
     argv = None
 
-    def __init__(self, parent, process = None, tty = False, wd = None, env = None, timeout = Timeout.default, level = 0, raw = True):
-        super(ssh_channel, self).__init__(timeout, level=level)
+    def __init__(self, parent, process = None, tty = False, wd = None, env = None, raw = True, *args, **kwargs):
+        super(ssh_channel, self).__init__(*args, **kwargs)
 
         # keep the parent from being garbage collected in some cases
         self.parent = parent
@@ -366,8 +366,8 @@ def libc(self):
                 return e
 
 class ssh_connecter(sock):
-    def __init__(self, parent, host, port, timeout = Timeout.default, level = None):
-        super(ssh_connecter, self).__init__(timeout, level = level)
+    def __init__(self, parent, host, port, *a, **kw):
+        super(ssh_connecter, self).__init__(*a, **kw)
 
         # keep the parent from being garbage collected in some cases
         self.parent = parent
@@ -398,8 +398,8 @@ def _close_msg(self):
 
 
 class ssh_listener(sock):
-    def __init__(self, parent, bind_address, port, timeout = Timeout.default, level = None):
-        super(ssh_listener, self).__init__(timeout, level = level)
+    def __init__(self, parent, bind_address, port, *a, **kw):
+        super(ssh_listener, self).__init__(*a, **kw)
 
         # keep the parent from being garbage collected in some cases
         self.parent = parent
@@ -478,8 +478,7 @@ class ssh(Timeout, Logger):
 
     def __init__(self, user, host, port = 22, password = None, key = None,
                  keyfile = None, proxy_command = None, proxy_sock = None,
-                 timeout = Timeout.default, level = None, cache = True,
-                 ssh_agent = False):
+                 level = None, cache = True, ssh_agent = False, *a, **kw):
         """Creates a new ssh connection.
 
         Arguments:
@@ -498,7 +497,7 @@ def __init__(self, user, host, port = 22, password = None, key = None,
 
         NOTE: The proxy_command and proxy_sock arguments is only available if a
         fairly new version of paramiko is used."""
-        super(ssh, self).__init__(timeout)
+        super(ssh, self).__init__(*a, **kw)
 
         Logger.__init__(self)
         if level is not None:
@@ -952,7 +951,7 @@ def which(self, program):
 
         return result
 
-    def system(self, process, tty = True, wd = None, env = None, timeout = Timeout.default, raw = True):
+    def system(self, process, tty = True, wd = None, env = None, timeout = None, raw = True):
         r"""system(process, tty = True, wd = None, env = None, timeout = Timeout.default, raw = True) -> ssh_channel
 
         Open a new channel with a specific process inside. If `tty` is True,
@@ -978,7 +977,10 @@ def system(self, process, tty = True, wd = None, env = None, timeout = Timeout.d
         if wd is None:
             wd = self.cwd
 
-        return ssh_channel(self, process, tty, wd, env, timeout, level = self.level, raw = raw)
+        if timeout is None:
+            timeout = self.timeout
+
+        return ssh_channel(self, process, tty, wd, env, timeout = timeout, level = self.level, raw = raw)
 
     #: Backward compatibility.  Use :meth:`system`
     run = system

pwnlib/tubes/tube.py

@@ -30,18 +30,18 @@ class tube(Timeout, Logger):
     #: and related functions.
     newline = '\n'
 
-    def __init__(self, timeout = default, level = None):
+    def __init__(self, timeout = default, level = None, *a, **kw):
         super(tube, self).__init__(timeout)
 
         Logger.__init__(self, None)
         if level is not None:
             self.setLevel(level)
 
-        self.buffer          = Buffer()
+        self.buffer = Buffer(*a, **kw)
         atexit.register(self.close)
 
     # Functions based on functions from subclasses
-    def recv(self, numb = 4096, timeout = default):
+    def recv(self, numb = None, timeout = default):
         r"""recv(numb = 4096, timeout = default) -> str
 
         Receives up to `numb` bytes of data from the tube, and returns
@@ -72,6 +72,7 @@ def recv(self, numb = 4096, timeout = default):
             [...] Received 0xc bytes:
                 'Hello, world'
         """
+        numb = self.buffer.get_fill_size(numb)
         return self._recv(numb, timeout) or ''
 
     def unrecv(self, data):
@@ -122,7 +123,7 @@ def _fillbuffer(self, timeout = default):
         data = ''
 
         with self.local(timeout):
-            data = self.recv_raw(4096)
+            data = self.recv_raw(self.buffer.get_fill_size())
 
         if data and self.isEnabledFor(logging.DEBUG):
             self.debug('Received %#x bytes:' % len(data))
@@ -141,12 +142,13 @@ def _fillbuffer(self, timeout = default):
         return data
 
 
-    def _recv(self, numb = 4096, timeout = default):
+    def _recv(self, numb = None, timeout = default):
         """_recv(numb = 4096, timeout = default) -> str
 
         Receives one chunk of from the internal buffer or from the OS if the
         buffer is empty.
         """
+        numb = self.buffer.get_fill_size(numb)
         data = ''
 
         # No buffered data, could not put anything in the buffer