1 files changed, 29 insertions, 16 deletions
diff --git a/sploit/payload/payload.py b/sploit/payload/payload.py
index b88e323..2a9521f 100644
--- a/sploit/payload/payload.py
+++ b/sploit/payload/payload.py
@@ -144,7 +144,9 @@ class Payload(IndexTbl):
 
     def __len__(self):
         """Return the size of the payload content in bytes."""
-        return len(bytes(self))
+        memo = {}
+        x = [ self.__lenof(i, memo) for i in range(len(self.__entries__)) ]
+        return sum(x)
 
     def __getindex__(self, index):
         """Return payload index value or address."""
@@ -167,7 +169,7 @@ class Payload(IndexTbl):
     # Payload helpers
 
     def __valueof(self, index, memo):
-        """Return a tuple: (addr of value, literal value) for index."""
+        """Return a tuple (addr of value, literal value) for index."""
         value = self.__entries__[index]
         addr = self.__addrof(index, memo)
         if isinstance(value, IndexEntry):
@@ -178,21 +180,27 @@ class Payload(IndexTbl):
     def __addrof(self, index, memo):
         """Return address (base + offset) for index."""
         index = self.__entries__.key2idx(index)
-        sizes = [ len(self.__bytesof(i, memo)) for i in range(index) ]
-        return self.base + sum(sizes)
-
-    def __bytesof(self, index, memo):
-        """Return byte output for index."""
         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
+            sizes = [ self.__lenof(i, memo) for i in range(index) ]
+            addr = self.base + sum(sizes)
+            memo[index] = addr
+            return addr
+
+    def __lenof(self, index, memo):
+        """Return element length for index."""
+        _, value = self.__valueof(index, memo)
+        if isinstance(value, PayloadEntry):
+            return value.payload_len(self)
+        return len(value)
+
+    def __bytesof(self, index, memo):
+        """Return byte output for index."""
+        _, value = self.__valueof(index, memo)
+        if isinstance(value, PayloadEntry):
+            return value.payload_bytes(self)
+        return bytes(value)
 
     def __prep_insertion(self, value, addr):
         """Initialize or type coerce input value for payload insert."""
@@ -206,6 +214,11 @@ class Payload(IndexTbl):
         elif type(value) is int:
             value = itob(value)
 
-        # Confirm value has a functional conversion to bytes
-        bytes(value)
+        try:
+            # Confirm value supports our required operations
+            len(value)
+            bytes(value)
+        except TypeError as ex:
+            raise TypeError(f"Payload: Bad type {type(value)} given") from ex
+
         return value