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."""