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>

2 files changed, 57 insertions, 33 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
diff --git a/sploit/payload/payload_entry.py b/sploit/payload/payload_entry.py
index 7088f83..4dca83d 100644
--- a/sploit/payload/payload_entry.py
+++ b/sploit/payload/payload_entry.py
@@ -20,15 +20,23 @@ class PayloadEntry(IndexEntry):
         """
         pass
 
+    def payload_len(self, payload):
+        """
+        Called to compute size of this entry.
+
+        Implement this method to calculate the length of this dynamic payload
+        entry.  self.base is set to the current entry address or offset.
+        """
+        raise NotImplementedError
+
     def payload_bytes(self, payload):
         """
         Called to generate bytes for this entry.
 
-        Override this method to generate and return the binary output for this
-        dynamic payload entry.  self.base is set to the current entry address
-        or offset.
+        Implement this method to generate the binary output for this dynamic
+        payload entry.  self.base is set to the current entry address or offset.
         """
-        return b""
+        raise NotImplementedError
 
 # Concrete payload entry definitions
 
@@ -43,6 +51,9 @@ class pointer(PayloadEntry):
             self.target = self.base
         self.target -= self.base
 
+    def payload_len(self, payload):
+        return arch.wordsize
+
     def payload_bytes(self, payload):
         return itob(self.target + self.base)
 
@@ -53,7 +64,11 @@ class padlen(PayloadEntry):
         self.size = size
         self.data = data
 
-    def _gen_padding(self, size):
+    def payload_len(self, payload):
+        return self.size - (self.base - payload.base)
+
+    def payload_bytes(self, payload):
+        size = self.payload_len(payload)
         data = self.data or arch.nopcode
         if size < 0:
             raise ValueError("padding: Available space is negative")
@@ -61,20 +76,17 @@ class padlen(PayloadEntry):
             raise ValueError("padding: Element does not divide the space evenly")
         return data * int(size)
 
-    def payload_bytes(self, payload):
-        return self._gen_padding(self.size - (self.base - payload.base))
-
 class padabs(padlen):
     """Generate padding to reach a target absolute address."""
 
-    def payload_bytes(self, payload):
-        return self._gen_padding(self.size - self.base)
+    def payload_len(self, payload):
+        return self.size - self.base
 
 class padrel(padlen):
     """Generate a fixed length of padding (aka: length relative to self)."""
 
-    def payload_bytes(self, payload):
-        return self._gen_padding(self.size)
+    def payload_len(self, payload):
+        return self.size
 
 class padalign(padlen):
     """Generate padding to reach next aligned address."""
@@ -83,9 +95,9 @@ class padalign(padlen):
         self.size = size
         self.data = data
 
-    def payload_bytes(self, payload):
+    def payload_len(self, payload):
         size = self.size or arch.alignment
-        return self._gen_padding(-self.base % size)
+        return -self.base % size
 
 class placeholder(padlen):
     """Generate fixed length of magic bytes, one word length by default."""
@@ -94,6 +106,5 @@ class placeholder(padlen):
         self.size = size
         self.data = _PLACEHOLDER_MAGIC
 
-    def payload_bytes(self, payload):
-        size = self.size or arch.wordsize
-        return self._gen_padding(size)
+    def payload_len(self, payload):
+        return self.size or arch.wordsize