Merge pull request #85 from trailofbits/func-call-args

Support injecting constant arguments in function call

fickling/fickle.py

@@ -517,6 +517,7 @@ def insert_magic_int(self, magic: int, index: int = -1):
     def insert_function_call_on_unpickled_object(
         self,
         function_definition: str,
+        constant_args: Optional[List[Any]] = None,
     ):
         """Insert and call a function that takes the unpickled object as parameter.
 
@@ -552,6 +553,11 @@ def insert_function_call_on_unpickled_object(
         self.insert(-1, Mark())
         self.insert(-1, Get.create(2))
 
+        # Add constant arguments
+        if constant_args:
+            for arg in constant_args:
+                self.insert(-1, ConstantOpcode.new(arg))
+
         # Now the stack contains [func, mark, model].
         # We need to add TUPLE which
         # packs the function arguments from the stack and then call REDUCE, which calls the injected
@@ -1429,6 +1435,9 @@ class ShortBinBytes(DynamicLength, ConstantOpcode):
     priority = Unicode.priority + 1
     length_bytes = 1
 
+    def encode_body(self) -> bytes:
+        return self.arg
+
     @classmethod
     def validate(cls, obj):
         if not isinstance(obj, bytes):