diff --git a/CHANGELOG.md b/CHANGELOG.md index cf91aeed1..fdc24761a 100644 --- a/CHANGELOG.md +++ b/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`) diff --git a/pwnlib/util/packing.py b/pwnlib/util/packing.py index 2f90c36bf..1a4fc9e19 100644 --- a/pwnlib/util/packing.py +++ b/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