Merge pull request #171 from trailofbits/fix-ast-unparse-crash

Fix ast.unparse() crash with malformed pickle files

Author: dhalf

fickling/fickle.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import keyword
 import marshal
 import re
 import struct
@@ -44,6 +45,63 @@ def is_std_module(module_name: str) -> bool:
     return in_stdlib(module_name) or module_name in BUILTIN_MODULE_NAMES
 
 
+def extract_identifier_from_ast_node(
+    node: ast.expr | str, fallback_prefix: str = "_malformed"
+) -> str:
+    """
+    Extract a valid Python identifier from an AST node.
+
+    For malformed pickle files where STACK_GLOBAL receives complex AST nodes
+    instead of strings, this function attempts to extract a meaningful identifier
+    or generates a safe fallback.
+
+    Args:
+        node: An AST expression node or string
+        fallback_prefix: Prefix for generated fallback identifiers
+
+    Returns:
+        A valid Python identifier string
+    """
+    # If already a string, validate and return or fix it
+    if isinstance(node, str):
+        if node.isidentifier() and not keyword.iskeyword(node):
+            return node
+        # For invalid string identifiers, create a safe fallback
+        node_hash = abs(hash(node)) % 100000
+        return f"{fallback_prefix}_str_{node_hash}"
+
+    # Handle ast.Name - extract the id attribute
+    if isinstance(node, ast.Name):
+        return node.id
+
+    # Handle ast.Attribute - return the attribute name
+    if isinstance(node, ast.Attribute):
+        return node.attr
+
+    # Handle ast.Call - extract identifier from the function being called
+    if isinstance(node, ast.Call):
+        return extract_identifier_from_ast_node(node.func, fallback_prefix)
+
+    # Handle ast.Constant - extract value if it's a valid identifier string
+    if isinstance(node, ast.Constant):
+        value = node.value
+        if isinstance(value, str):
+            if value.isidentifier() and not keyword.iskeyword(value):
+                return value
+            # For invalid constant strings, create a deterministic fallback
+            value_hash = abs(hash(value)) % 100000
+            return f"{fallback_prefix}_const_{value_hash}"
+        # For non-string constants (int, float, etc.), create a type-based fallback
+        type_name = type(value).__name__
+        value_hash = abs(hash(value)) % 100000
+        return f"{fallback_prefix}_{type_name}_{value_hash}"
+
+    # For any other complex node types, generate a safe placeholder
+    node_type = type(node).__name__.lower()
+    node_hash = abs(hash(id(node))) % 100000
+    return f"{fallback_prefix}_{node_type}_{node_hash}"
+
+
 class MarkObject:
     pass
 
@@ -1060,20 +1118,45 @@ class StackGlobal(NoOp):
     def run(self, interpreter: Interpreter):
         attr = interpreter.stack.pop()
         module = interpreter.stack.pop()
+
+        # Extract values from ast.Constant nodes
         if isinstance(module, ast.Constant):
             module = module.value
         if isinstance(attr, ast.Constant):
             attr = attr.value
 
-        # normalize module and attr to strings
-        if not isinstance(module, str) or not isinstance(attr, str):
+        # Normalize module and attr to strings, extracting meaningful identifiers from AST nodes
+        module_needs_extraction = not isinstance(module, str)
+        attr_needs_extraction = not isinstance(attr, str)
+
+        if module_needs_extraction or attr_needs_extraction:
             sys.stdout.write(
                 f"Warning: malformed pickle file. STACK_GLOBAL called with invalid types. "
                 f"'Module' is {type(module).__name__} ({module!r}), 'Attr' is {type(attr).__name__} ({attr!r}). "
-                f"Expected str; casting to string to continue analysis.\n"
+                f"Expected str; extracting identifiers to continue analysis.\n"
             )
-            module = str(module)
-            attr = str(attr)
+
+            if module_needs_extraction:
+                module = extract_identifier_from_ast_node(
+                    module, fallback_prefix="_malformed_module"
+                )
+            if attr_needs_extraction:
+                attr = extract_identifier_from_ast_node(attr, fallback_prefix="_malformed_attr")
+
+        # Final validation: ensure both are valid identifier strings
+        if not isinstance(module, str) or not isinstance(attr, str):
+            raise TypeError(
+                f"Failed to extract valid identifiers from STACK_GLOBAL arguments. "
+                f"Module: {type(module).__name__}, Attr: {type(attr).__name__}"
+            )
+
+        if not module.isidentifier() or not attr.isidentifier():
+            raise ValueError(
+                f"Extracted identifiers are not valid Python identifiers. "
+                f"Module: {module!r}, Attr: {attr!r}"
+            )
+
+        # Continue with normal processing
         if module in ("__builtin__", "__builtins__", "builtins"):
             # no need to emit an import for builtins!
             pass
@@ -1133,6 +1216,10 @@ class BinPut(Opcode):
     def run(self, interpreter: Interpreter):
         interpreter.memory[self.arg] = interpreter.stack[-1]
 
+    def encode_body(self):
+        assert self.arg <= 255, "BINPUT only supports 1-byte memo indexing"
+        return bytes([self.arg])
+
 
 class LongBinPut(BinPut):
     name = "LONG_BINPUT"
@@ -1422,6 +1509,10 @@ def run(self, interpreter: Interpreter):
         else:
             interpreter.stack.append(interpreter.memory[self.arg])
 
+    def encode_body(self):
+        assert self.arg <= 255, "BINGET only supports 1-byte memo indexing"
+        return bytes([self.arg])
+
 
 class LongBinGet(Opcode):
     name = "LONG_BINGET"

test/test_crashes.py

@@ -8,7 +8,23 @@
 from functools import wraps
 from unittest import TestCase
 
-from fickling.fickle import Pickled
+from fickling.fickle import (
+    BinGet,
+    BinInt1,
+    BinPut,
+    BinUnicode,
+    Global,
+    Mark,
+    Memoize,
+    Pickled,
+    Pop,
+    Proto,
+    Reduce,
+    ShortBinUnicode,
+    StackGlobal,
+    Stop,
+    Tuple,
+)
 
 
 def unparse_test(pickled: bytes):
@@ -60,3 +76,74 @@ def test_pop_mark(self):
     def test_obj(self):
         """Tests the correctness of the OBJ opcode"""
         pass
+
+    # Based on the CTF challenge shared in https://github.com/trailofbits/fickling/issues/125.
+    def test_stack_global_dynamic_import(self):
+        alphabet = (
+            "Jw~[v5QpA(BY%aKnyT&*x0r9-OpfF}HN4$GU2VhS@XEq!Zt>6_R7#]1b{z3M^D?)d8eImgckPLiuoClW<js"
+        )
+        pickled = Pickled(
+            [
+                Proto.create(4),
+                # Save itemgetter to memo[0]
+                Global.create("operator", "itemgetter"),
+                BinPut(0),
+                Pop(),
+                # _var0 = getattr('', 'join')
+                Global.create("builtins", "getattr"),
+                Mark(),
+                ShortBinUnicode(""),
+                ShortBinUnicode("join"),
+                Tuple(),
+                Reduce(),
+                Memoize(),  # memo[1] = join method
+                # _var1 = itemgetter(77, 83)(ALPHABET) -> ('o', 's')
+                BinGet(0),
+                Mark(),
+                BinInt1(77),
+                BinInt1(83),
+                Tuple(),
+                Reduce(),  # itemgetter(77, 83)
+                Mark(),
+                BinUnicode(alphabet),
+                Tuple(),
+                Reduce(),
+                Memoize(),  # memo[2] = ('o', 's')
+                # _var2 = ''.join(_var1) -> "os"
+                BinGet(1),
+                Mark(),
+                BinGet(2),
+                Tuple(),
+                Reduce(),
+                Memoize(),  # memo[3] = "os"
+                # _var3 = itemgetter(83, 16, 83, 47, 67, 69)(ALPHABET) -> ('s','y','s','t','e','m')
+                BinGet(0),
+                Mark(),
+                BinInt1(83),
+                BinInt1(16),
+                BinInt1(83),
+                BinInt1(47),
+                BinInt1(67),
+                BinInt1(69),
+                Tuple(),
+                Reduce(),  # itemgetter(...)
+                Mark(),
+                BinUnicode(alphabet),
+                Tuple(),
+                Reduce(),
+                Memoize(),  # memo[4] = ('s','y','s','t','e','m')
+                # _var4 = ''.join(_var3) -> "system"
+                BinGet(1),
+                Mark(),
+                BinGet(4),
+                Tuple(),
+                Reduce(),
+                Memoize(),  # memo[5] = "system"
+                # from _var2 import _var4 (dynamic import via StackGlobal)
+                BinGet(3),  # "os"
+                BinGet(5),  # "system"
+                StackGlobal(),
+                Stop(),
+            ]
+        )
+        unparse(pickled.ast)