payload: Separate length and bytes calculations

Previously, the len(payload) operation required the generation of the
full payload binary content, in order to count how many bytes long it
was.  This is no longer the case, as there are opportunities for
optimizations, primarily regarding fixed-length dynamic payload entries
where we can simply grab the size parameter without having to generate a
buffer.

In addition to potential speedups, this fix also allows the user to
insert PayloadEntry pointers for fields which are not yet present in the
payload being built (ie: whenever the pointer is to exist before the
pointed-to data).  Whereas previously, the inability to generate the
ill-formed pointer would break length calculations necessary to insert
additional data.

Signed-off-by: Malfurious <m@lfurio.us>

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