diff options
| -rw-r--r-- | sploit/payload/payload.py | 36 |
1 files changed, 22 insertions, 14 deletions
diff --git a/sploit/payload/payload.py b/sploit/payload/payload.py index a59eed2..b88e323 100644 --- a/sploit/payload/payload.py +++ b/sploit/payload/payload.py @@ -84,6 +84,7 @@ class Payload(IndexTbl): """Return human-readable Payload.""" FMT = "\n{:<20} {:<20} {:<20}" s = f"{len(self.__entries__)} items, {len(self)} bytes @ {hex(self)}" + memo = {} if len(self.__entries__) > 0: s += FMT.format("ADDRESS", "SYMBOL", "DATA") @@ -93,8 +94,8 @@ class Payload(IndexTbl): key = "(unkeyed)" if key is None else str(key) key = f"[{key}]" if isinstance(value, IndexEntry) else key - addr = self.__addrof(i) - data = str(self.__bytesof(i)) + addr = self.__addrof(i, memo) + data = str(self.__bytesof(i, memo)) if len(data) > _REPR_DATA_LEN: data = data[:_REPR_DATA_LEN] + " ..." @@ -104,7 +105,8 @@ class Payload(IndexTbl): def __bytes__(self): """Return calculated payload bytes.""" - x = [ self.__bytesof(i) for i in range(len(self.__entries__)) ] + memo = {} + x = [ self.__bytesof(i, memo) for i in range(len(self.__entries__)) ] return b"".join(x) def __call__(self, *args): @@ -146,13 +148,13 @@ class Payload(IndexTbl): def __getindex__(self, index): """Return payload index value or address.""" - value, _ = self.__valueof(index) + value, _ = self.__valueof(index, {}) return value def __setindex__(self, index, value): """Set payload index value.""" try: - addr = self.__addrof(index) + addr = self.__addrof(index, {}) except KeyError: addr = self.end() value = self.__prep_insertion(value, addr) @@ -164,27 +166,33 @@ class Payload(IndexTbl): # Payload helpers - def __valueof(self, index): + def __valueof(self, index, memo): """Return a tuple: (addr of value, literal value) for index.""" value = self.__entries__[index] - addr = self.__addrof(index) + addr = self.__addrof(index, memo) if isinstance(value, IndexEntry): value @= addr return value, value return addr, value - def __addrof(self, index): + def __addrof(self, index, memo): """Return address (base + offset) for index.""" index = self.__entries__.key2idx(index) - sizes = [ len(self.__bytesof(i)) for i in range(index) ] + sizes = [ len(self.__bytesof(i, memo)) for i in range(index) ] return self.base + sum(sizes) - def __bytesof(self, index): + def __bytesof(self, index, memo): """Return byte output for index.""" - _, value = self.__valueof(index) - if isinstance(value, PayloadEntry): - return value.payload_bytes(self) - return bytes(value) + try: + return memo[index] + except KeyError: + _, value = self.__valueof(index, memo) + if isinstance(value, PayloadEntry): + data = value.payload_bytes(self) + else: + data = bytes(value) + memo[index] = data + return data def __prep_insertion(self, value, addr): """Initialize or type coerce input value for payload insert.""" |