1#! /usr/bin/env python3
2
3"""
4Compare checksums for wheels in :mod:`ensurepip` against the Cheeseshop.
5
6When GitHub Actions executes the script, output is formatted accordingly.
7https://docs.github.com/en/actions/using-workflows/workflow-commands-for-github-actions#setting-a-notice-message
8"""
9
10import hashlib
11import json
12import os
13import re
14from pathlib import Path
15from urllib.request import urlopen
16
17PACKAGE_NAMES = ("pip", "setuptools")
18ENSURE_PIP_ROOT = Path(__file__).parent.parent.parent / "Lib/ensurepip"
19WHEEL_DIR = ENSURE_PIP_ROOT / "_bundled"
20ENSURE_PIP_INIT_PY_TEXT = (ENSURE_PIP_ROOT / "__init__.py").read_text(encoding="utf-8")
21GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
22
23
24def print_notice(file_path: str, message: str) -> None:
25    if GITHUB_ACTIONS:
26        message = f"::notice file={file_path}::{message}"
27    print(message, end="\n\n")
28
29
30def print_error(file_path: str, message: str) -> None:
31    if GITHUB_ACTIONS:
32        message = f"::error file={file_path}::{message}"
33    print(message, end="\n\n")
34
35
36def verify_wheel(package_name: str) -> bool:
37    # Find the package on disk
38    package_path = next(WHEEL_DIR.glob(f"{package_name}*.whl"), None)
39    if not package_path:
40        print_error("", f"Could not find a {package_name} wheel on disk.")
41        return False
42
43    print(f"Verifying checksum for {package_path}.")
44
45    # Find the version of the package used by ensurepip
46    package_version_match = re.search(
47        f'_{package_name.upper()}_VERSION = "([^"]+)', ENSURE_PIP_INIT_PY_TEXT
48    )
49    if not package_version_match:
50        print_error(
51            package_path,
52            f"No {package_name} version found in Lib/ensurepip/__init__.py.",
53        )
54        return False
55    package_version = package_version_match[1]
56
57    # Get the SHA 256 digest from the Cheeseshop
58    try:
59        raw_text = urlopen(f"https://pypi.org/pypi/{package_name}/json").read()
60    except (OSError, ValueError):
61        print_error(package_path, f"Could not fetch JSON metadata for {package_name}.")
62        return False
63
64    release_files = json.loads(raw_text)["releases"][package_version]
65    for release_info in release_files:
66        if package_path.name != release_info["filename"]:
67            continue
68        expected_digest = release_info["digests"].get("sha256", "")
69        break
70    else:
71        print_error(package_path, f"No digest for {package_name} found from PyPI.")
72        return False
73
74    # Compute the SHA 256 digest of the wheel on disk
75    actual_digest = hashlib.sha256(package_path.read_bytes()).hexdigest()
76
77    print(f"Expected digest: {expected_digest}")
78    print(f"Actual digest:   {actual_digest}")
79
80    if actual_digest != expected_digest:
81        print_error(
82            package_path, f"Failed to verify the checksum of the {package_name} wheel."
83        )
84        return False
85
86    print_notice(
87        package_path,
88        f"Successfully verified the checksum of the {package_name} wheel.",
89    )
90    return True
91
92
93if __name__ == "__main__":
94    exit_status = 0
95    for package_name in PACKAGE_NAMES:
96        if not verify_wheel(package_name):
97            exit_status = 1
98    raise SystemExit(exit_status)
99