Repository: Gallopsled/pwntools

Issue:
fit() fails when keys are not ordered
This results in what appears to be an infinite loop.

```
from pwn import *

fit({
	0x61616168: 1,
	0x61616167: 2
})
```

This should be valid since these correspond to consecutive cyclic values:

```
>>> pack(0x61616167)
'gaaa'
>>> pack(0x61616168)
'haaa'
```

These also work individually, and do not overlap:

```
>>> fit({
... 0x61616167: 2,
... })
'aaaabaaacaaadaaaeaaafaaa\x02\x00\x00\x00'
>>> fit({
... 0x61616168: 1,
... })
'aaaabaaacaaadaaaeaaafaaagaaa\x01\x00\x00\x00'
```


Discussion:
I can confirm this works as expected on 3.9.3:

```
>>> pwnlib.version
'3.9.3'
>>> fit({
... 0x61616168: 1,
... 0x61616167: 2
... })
'aaaabaaacaaadaaaeaaafaaa\x02\x00\x00\x00\x01\x00\x00\x00'
```
I've confirmed this comes down to ordering of the fields.  On `dev`, I see:

```
>>> fit({0x61616167: 'a', 0x61616168: 'b'})
b'aaaabaaacaaadaaaeaaafaaaaaaab'
>>> fit({0x61616168: 'a', 0x61616167: 'b'})
^C (infinite loop)
```
The fix I've uploaded fixes this specific issue, but there are likely more issues hiding under the surface when switching between e.g. `{0: ...}` and `{0x61616116: ...}`
This apparently only fixes the issue on Python3, which ensures sorted-ness of dictionary items.  This requires an add'l fix for Python2.

Resolving patch:
diff --git a/pwnlib/util/packing.py b/pwnlib/util/packing.py
--- a/pwnlib/util/packing.py
+++ b/pwnlib/util/packing.py
@@ -492,9 +492,11 @@ def _fit(pieces, preprocessor, packer, filler):
     pad = bytearray()
     def fill(key):
         key = bytearray(key)
-        while len(pad) < len(key) or pad[-len(key):] != key:
+        offset = pad.find(key)
+        while offset == -1:
             pad.append(next(filler))
-        return len(pad) - len(key)
+            offset = pad.find(key, -len(key))
+        return offset
 
     # Key conversion:
     # - convert str/unicode keys to offsets
@@ -549,7 +551,6 @@ def _flat(args, preprocessor, packer, filler):
         elif isinstance(arg, (list, tuple)):
             val = _flat(arg, preprocessor, packer, filler)
         elif isinstance(arg, dict):
-            arg = collections.OrderedDict(((k,v) for k,v in sorted(arg.items())))
             filler, val = _fit(arg, preprocessor, packer, filler)
         elif isinstance(arg, bytes):
             val = arg
@@ -617,37 +618,100 @@ def flat(*args, **kwargs):
 
     Examples:
 
-      >>> context.clear()
-      >>> flat(1, "test", [[["AB"]*2]*3], endianness = 'little', word_size = 16, sign = False)
-      b'\x01\x00testABABABABABAB'
-      >>> flat([1, [2, 3]], preprocessor = lambda x: str(x+1))
-      b'234'
-      >>> flat({12: 0x41414141,
-      ...       24: 'Hello',
-      ...      })
-      b'aaaabaaacaaaAAAAeaaafaaaHello'
-      >>> flat({'caaa': ''})
-      b'aaaabaaa'
-      >>> flat({12: 'XXXX'}, filler = (ord('A'), ord('B')), length = 20)
-      b'ABABABABABABXXXXABAB'
-      >>> flat({ 8: [0x41414141, 0x42424242],
-      ...       20: 'CCCC'})
-      b'aaaabaaaAAAABBBBeaaaCCCC'
-      >>> flat({ 0x61616162: 'X'})
-      b'aaaaX'
-      >>> flat({4: {0: 'X', 4: 'Y'}})
-      b'aaaaXaaaY'
-      >>> flat({0x61616161:'x', 0x61616162:'y'})
-      b'xaaay'
-      >>> flat({0x61616162:'y', 0x61616161:'x'})
-      b'xaaay'
-      >>> fit({
-      ...     0x61616161: 'a',
-      ...     1: 'b',
-      ...     0x61616161+2: 'c',
-      ...     3: 'd',
-      ... })
-      b'abadbaaac'
+        (Test setup, please ignore)
+    
+        >>> context.clear()
+
+        Basic usage of :meth:`flat` works similar to the pack() routines.
+
+        >>> flat(4)
+        b'\x04\x00\x00\x00'
+
+        :meth:`flat` works with strings, bytes, lists, and dictionaries.
+
+        >>> flat(b'X')
+        b'X'
+        >>> flat([1,2,3])
+        b'\x01\x00\x00\x00\x02\x00\x00\x00\x03\x00\x00\x00'
+        >>> flat({4:'X'})
+        b'aaaaX'
+
+        :meth:`.flat` flattens all of the values provided, and allows nested lists
+        and dictionaries.
+
+        >>> flat([{4:'X'}] * 2)
+        b'aaaaXaaacX'
+        >>> flat([[[[[[[[[1]]]], 2]]]]])
+        b'\x01\x00\x00\x00\x02\x00\x00\x00'
+
+        You can also provide additional arguments like endianness, word-size, and
+        whether the values are treated as signed or not.
+
+        >>> flat(1, "test", [[["AB"]*2]*3], endianness = 'little', word_size = 16, sign = False)
+        b'\x01\x00testABABABABABAB'
+
+        A preprocessor function can be provided in order to modify the values in-flight.
+        This example converts increments each value by 1, then converts to a string.
+
+        >>> flat([1, [2, 3]], preprocessor = lambda x: str(x+1))
+        b'234'
+
+        Using dictionaries is a fast way to get specific values at specific offsets,
+        without having to do ``data += "foo"`` repeatedly.
+
+        >>> flat({12: 0x41414141,
+        ...       24: 'Hello',
+        ...      })
+        b'aaaabaaacaaaAAAAeaaafaaaHello'
+
+        Dictionary usage permits directly using values derived from :func:`.cyclic`.
+        See :func:`.cyclic`, :function:`pwnlib.context.context.cyclic_alphabet`, and :data:`.context.cyclic_size`
+        for more options.  
+
+        The cyclic pattern can be provided as either the text or hexadecimal offset.
+
+        >>> flat({ 0x61616162: 'X'})
+        b'aaaaX'
+        >>> flat({'baaa': 'X'})
+        b'aaaaX'
+
+        Fields do not have to be in linear order, and can be freely mixed.
+        This also works with cyclic offsets.
+
+        >>> flat({2: 'A', 0:'B'})
+        b'BaA'
+        >>> flat({0x61616161:'x', 0x61616162:'y'})
+        b'xaaay'
+        >>> flat({0x61616162:'y', 0x61616161:'x'})
+        b'xaaay'
+
+        Fields do not have to be in order, and can be freely mixed.
+
+        >>> flat({'caaa': 'XXXX', 16: '\x41', 20: 0xdeadbeef})
+        b'aaaabaaaXXXXdaaaAaaa\xef\xbe\xad\xde'
+        >>> flat({ 8: [0x41414141, 0x42424242], 20: 'CCCC'})
+        b'aaaabaaaAAAABBBBeaaaCCCC'
+        >>> fit({
+        ...     0x61616161: 'a',
+        ...     1: 'b',
+        ...     0x61616161+2: 'c',
+        ...     3: 'd',
+        ... })
+        b'abadbaaac'
+
+        By default, gaps in the data are filled in with the :meth:`.cyclic` pattern.
+        You can customize this by providing an iterable or method for the ``filler``
+        argument.
+
+        >>> flat({12: 'XXXX'}, filler = (ord('A'), ord('B')), length = 20)
+        b'ABABABABABABXXXXABAB'
+
+        Nested dictionaries also work as expected.
+
+        >>> flat({4: {0: 'X', 4: 'Y'}})
+        b'aaaaXaaaY'
+        >>> fit({4: {4: 'XXXX'}})
+        b'aaaabaaaXXXX'
     """
     # HACK: To avoid circular imports we need to delay the import of `cyclic`
     from pwnlib.util import cyclic
@@ -670,7 +734,7 @@ def flat(*args, **kwargs):
     return out
 
 def fit(*args, **kwargs):
-    """Legacy alias for `:func:flat`"""
+    """Legacy alias for :func:`flat`"""
     return flat(*args, **kwargs)
 
 """
diff --git a/pwnlib/version.py b/pwnlib/version.py
--- a/pwnlib/version.py
+++ b/pwnlib/version.py
@@ -1 +1 @@
-__version__ = '4.1.3'
+__version__ = '4.1.4'
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.3',
+    version              = '4.1.4',
     data_files           = [('',
                              glob.glob('*.md') + glob.glob('*.txt')),
                             ],


Test patch: