Permit negative indices in fit() and flat()

CHANGELOG.md

@@ -65,6 +65,7 @@ To be released on Jun 30, 2020.
 - [#1605][1605] Add to `fiddling.hexdump` a way to suppress the total at the end
 - [#1613][1613] Permit `--password` for `pwn template`
 - [#1564][1564] Fix `asm()` and `disasm()` for PowerPC64, MIPS64, Sparc64
+- [#1621][1621] Permit negative values in flat() and fit()
 
 [1576]: https://github.com/Gallopsled/pwntools/pull/1576
 [1584]: https://github.com/Gallopsled/pwntools/pull/1584
@@ -74,6 +75,7 @@ To be released on Jun 30, 2020.
 [1605]: https://github.com/Gallopsled/pwntools/pull/1605
 [1613]: https://github.com/Gallopsled/pwntools/pull/1613
 [1564]: https://github.com/Gallopsled/pwntools/pull/1564
+[1621]: https://github.com/Gallopsled/pwntools/pull/1621
 
 ## 4.2.0 (`beta`)
 

pwnlib/util/packing.py

@@ -42,9 +42,12 @@
 
 from pwnlib.context import LocalNoarchContext
 from pwnlib.context import context
+from pwnlib.log import getLogger
+
 from pwnlib.util import iters
 
 mod = sys.modules[__name__]
+log = getLogger(__name__)
 
 def pack(number, word_size = None, endianness = None, sign = None, **kwargs):
     """pack(number, word_size = None, endianness = None, sign = None, **kwargs) -> str
@@ -487,6 +490,7 @@ def make_unpacker(word_size = None, endianness = None, sign = None, **kwargs):
         return lambda number: unpack(number, word_size, endianness, sign)
 
 def _fit(pieces, preprocessor, packer, filler):
+
     # Pulls bytes from `filler` and adds them to `pad` until it ends in `key`.
     # Returns the index of `key` in `pad`.
     pad = bytearray()
@@ -524,6 +528,14 @@ def fill(key):
 
     # Build output
     out = b''
+
+    # Negative indices need to be removed and then re-submitted
+    negative = {k:v for k,v in pieces.items() if isinstance(k, int) and k<0}
+
+    for k in negative:
+        del pieces[k]
+
+    # Positive output
     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)))
@@ -535,7 +547,24 @@ def fill(key):
         # Recursively flatten data
         out += _flat([v], preprocessor, packer, filler)
 
-    return filler, out
+    # Now do negative indices
+    out_negative = b''
+    if negative:
+        most_negative = min(negative.keys())
+        for k, v in sorted(negative.items()):
+            k += -most_negative
+
+            if k < len(out_negative):
+                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_negative) < k:
+                out_negative += p8(next(filler))
+
+            # Recursively flatten data
+            out_negative += _flat([v], preprocessor, packer, filler)
+
+    return filler, out_negative + out
 
 def _flat(args, preprocessor, packer, filler):
     out = []
@@ -714,6 +743,12 @@ def flat(*args, **kwargs):
         b'aaaaXaaaY'
         >>> fit({4: {4: 'XXXX'}})
         b'aaaabaaaXXXX'
+
+        Negative indices are also supported, though this only works for integer
+        keys.
+
+        >>> flat({-4: 'x', -1: 'A', 0: '0', 4:'y'})
+        b'xaaA0aaay'
     """
     # HACK: To avoid circular imports we need to delay the import of `cyclic`
     from pwnlib.util import cyclic