Repository: Gallopsled/pwntools

Issue:
fit() and flat() broken for simple uses of filler=
The `flat` and `fit` routines take an arbitrary iterable for the `filler=` argument.  This seems to be broken in recent releases.

# Pwntools 3.13.0

```
>>> fit({10: 'A'}, length=32, filler='_')
'__________A_____________________'
```

# Pwntools 4.2.0 (works)

```
>>> fit({10: 'A'}, length=32, filler=[ord('_')])
b'__________A_____________________'
```

# Pwntools 4.2.0 (broken, but should work)

```
>>> fit({10: 'A'}, length=32, filler='_')
---------------------------------------------------------------------------
error                                     Traceback (most recent call last)
<ipython-input-1-304fe6d91d58> in <module>()
----> 1 fit({10: 'A'}, length=32, filler='_')

/home/pwntools/pwntools/pwnlib/util/packing.py in fit(*args, **kwargs)
    658 def fit(*args, **kwargs):
    659     """Legacy alias for `:func:flat`"""
--> 660     return flat(*args, **kwargs)
    661
    662 """

/home/pwntools/pwntools/pwnlib/context/__init__.py in setter(*a, **kw)
   1457         kw.setdefault('arch', 'none')
   1458         with context.local(**{k:kw.pop(k) for k,v in tuple(kw.items()) if isinstance(getattr(ContextType, k, None), property)}):
-> 1459             return function(*a, **kw)
   1460     return setter
   1461

/home/pwntools/pwntools/pwnlib/util/packing.py in flat(*args, **kwargs)
    647
    648     filler = iters.cycle(filler)
--> 649     out = _flat(args, preprocessor, make_packer(), filler)
    650
    651     if length:

/home/pwntools/pwntools/pwnlib/util/packing.py in _flat(args, preprocessor, packer, filler)
    549             val = _flat(arg, preprocessor, packer, filler)
    550         elif isinstance(arg, dict):
--> 551             filler, val = _fit(arg, preprocessor, packer, filler)
    552         elif isinstance(arg, bytes):
    553             val = arg

/home/pwntools/pwntools/pwnlib/util/packing.py in _fit(pieces, preprocessor, packer, filler)
    528         # Fill up to offset
    529         while len(out) < k:
--> 530             out += p8(next(filler))
    531
    532         # Recursively flatten data

/home/pwntools/pwntools/pwnlib/context/__init__.py in setter(*a, **kw)
   1457         kw.setdefault('arch', 'none')
   1458         with context.local(**{k:kw.pop(k) for k,v in tuple(kw.items()) if isinstance(getattr(ContextType, k, None), property)}):
-> 1459             return function(*a, **kw)
   1460     return setter
   1461

/home/pwntools/pwntools/pwnlib/util/packing.py in routine(number)
    338                 ("little", False):  lu,
    339                 ("big",    True  ): bs,
--> 340                 ("big",    False):  bu}[endian, signed](number)
    341
    342     routine.__name__ = name

/home/pwntools/pwntools/pwnlib/util/packing.py in routine(data)
    303
    304     def routine(data):
--> 305         return ops[op](fmt,data)
    306     routine.__name__ = routine.__qualname__ = name
    307

error: cannot convert argument to integer
```


Discussion:
Following up, it looks like there *were* tests for this in Pwntools 3.13, but they were rewritten at some point.  Pwntools 3.13 shows:

```
      >>> fit({12: 'XXXX'}, filler = 'AB', length = 20)
      'ABABABABABABXXXXABAB'
```
It looks like these are the changes that broke things.

```diff
@@ -516,14 +520,14 @@ def _fit(pieces, preprocessor, packer, filler):
     filler = iters.chain(pad, filler)

     # Build output
-    out = ''
+    out = b''
     for k, v in sorted(pieces.items()):
         if k < len(out):
             raise ValueError("flat(): data at offset %d overlaps with previous data which ends at offset %d" % (k, len(out)))

         # Fill up to offset
         while len(out) < k:
-            out += filler.next()
+            out += p8(next(filler))
```

and

```diff
@@ -647,7 +651,7 @@ def flat(*args, **kwargs):
     if length:
         if len(out) > length:
             raise ValueError("flat(): Arguments does not fit within `length` (= %d) bytes" % length)
-        out += ''.join(filler.next() for _ in xrange(length - len(out)))
+        out += b''.join(p8(next(filler)) for _ in range(length - len(out)))
```

Removing the `p8` doesn't fix the issue, but seems to be on the right track
It looks like another part of the issue is the internal switch to use `bytearray`?  I'm not sure, still trying to diagnose this
It looks like something underlying `iters.cycle` also changed significantly, which comes from `itertools.cycle`.

It seems the `p8` above were added to handle this situation,  which makes things sort-of work on Python3, and broken on Python2.  This seems to be a core Py3-vs-Py2 issue that needs to be worked around.

# Python2

```
>>> x = iters.cycle(b'asdf')
>>> next(x)
'a'
>>> next(x)
's'
>>> next(x)
'd'
>>> next(x)
'f'
```

# Python3

### With bytes

```
>>> x = iters.cycle(b'asdf')
>>> next(x)
97
>>> next(x)
115
>>> next(x)
100
>>> next(x)
102
```

### With a string

```
>>> x = iters.cycle('asdf')
>>> next(x)
'a'
>>> next(x)
's'
>>> next(x)
'd'
>>> next(x)
'f'
```


Honestly this is a Python3 gaffe and I think the easiest way to work around it is to hack the `p8` routine to do the "magical correct thing" if the argument is a `bytes`-like object.
> It looks like something underlying `iters.cycle` also changed significantly, which comes from `itertools.cycle`.

@zachriggle Well, just this `iters.cycle` 'bug' is a total WONTFIX, since this function has little right to object to python iteration mechanism. It would just not be consistent to treat some types special instead of just giving the impression of `chain(iter(x), iter(x), iter(x), ...)` uniformly for all types. Whatever a type yields while iterating, is yielded by cycling it.

As for the actual issue, though, `p8` was just a choice of the moment (totally mea culpa), should be definitely changed to something more generic, like, say, `flat(x)[0]`? I'm open for ideas, but the following is to be considered.

So what if:
* The iterable is yielding strings longer than 2? Could be resolved by something like `lazy_flat(cycle(filler))`, which would return a generator of uint8s or just unsupported.
* The iterable is malicious, yielding empty strings all the time? Could be stated that this is not supported, or special-casing empty iterables / of known length 0 (before passing to `cycle()`? does `next(cycle(''))` hang forever?).
* The iterable yields integers not in range(256)? Could be just unsupported or encoded like in bullet no. 1 (if so, what with endianness?) --- generally seems like a mess that needs special-casing and different treatment for <256 and the rest (we want to still support passing `filler=b'_'`, right?  ... right...?).
* (possible more useful edge-case scenarios, but I did not think more).

~But hey, we could also make a cute hack that makes p8 be like context._encode on strings!~
EDIT: Actually not a chance I am personally implementing such a hack, although it's so tempting for a user to have a black box do things for them, isn't it.
Ultimately the issue comes down to:

## Python3
```py
>>> x = b'a'
>>> x[0] == x
False
>>> x[0]
97
```

## Python2

```py
>>> x = b'a'
>>> x[0] == x
True
>>> x[0]
'a'
```

So it looks like iterating over a byte-array is generating ints, in which case the `p8` is necessary on Python3 but not necessary on Python2.

I'll spend some time seeing if I can come up with a more elegant solution or wrapper for this.
It also occurs to me that using `flat` is a good way to ensure we always get a `bytes`-like object regardless of what is passed in, and might be a good thing to start using from e.g. `tubes` instead of `.encode()` and `.decode()`.

Resolving patch:
diff --git a/pwnlib/util/packing.py b/pwnlib/util/packing.py
--- a/pwnlib/util/packing.py
+++ b/pwnlib/util/packing.py
@@ -703,7 +703,9 @@ def flat(*args, **kwargs):
         You can customize this by providing an iterable or method for the ``filler``
         argument.
 
-        >>> flat({12: 'XXXX'}, filler = (ord('A'), ord('B')), length = 20)
+        >>> flat({12: 'XXXX'}, filler = b'_', length = 20)
+        b'____________XXXX____'
+        >>> flat({12: 'XXXX'}, filler = b'AB', length = 20)
         b'ABABABABABABXXXXABAB'
 
         Nested dictionaries also work as expected.
@@ -720,6 +722,9 @@ def flat(*args, **kwargs):
     filler       = kwargs.pop('filler', cyclic.de_bruijn())
     length       = kwargs.pop('length', None)
 
+    if isinstance(filler, str):
+        filler = bytearray(six.ensure_binary(filler))
+
     if kwargs != {}:
         raise TypeError("flat() does not support argument %r" % kwargs.popitem()[0])
 
diff --git a/pwnlib/version.py b/pwnlib/version.py
--- a/pwnlib/version.py
+++ b/pwnlib/version.py
@@ -1 +1 @@
-__version__ = '4.1.4'
+__version__ = '4.1.5'
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -83,7 +83,7 @@
     name                 = 'pwntools',
     python_requires      = '>=2.7',
     packages             = find_packages(),
-    version              = '4.1.4',
+    version              = '4.1.5',
     data_files           = [('',
                              glob.glob('*.md') + glob.glob('*.txt')),
                             ],


Test patch: