cli: Support friendlier syntax for verify pypi command (#87)

facutuesca · web-flow · commit 86cc7de8a503 · 2025-02-11T11:05:25.000-05:00
* cli: Support friendlier syntax for `verify pypi` command

Signed-off-by: Facundo Tuesca &lt;facundo.tuesca@trailofbits.com&gt;

* fixup! cli: Support friendlier syntax for `verify pypi` command

---------

Signed-off-by: Facundo Tuesca &lt;facundo.tuesca@trailofbits.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - The CLI subcommand `verify attestation` now supports `.slsa.attestation`
   files. When verifying an artifact, both `.publish.attestation` and
   `.slsa.attestation` files are used (if present).
+- The CLI subcommand `verify pypi` now supports a friendlier
+  syntax to specify the artifact to verify. The artifact can now be
+  specified with a `pypi:` prefix followed by the filename, e.g:
+  `pypi:sampleproject-1.0.0.tar.gz`. The old way (passing
+  the direct URL) is still supported.
 
 ## [0.0.21]
 
diff --git a/README.md b/README.md
@@ -142,19 +142,21 @@ pypi-attestations verify attestation  \
 
 ### Verifying a PyPI package
 > [!NOTE]
-> The URL must be a direct link to the distribution artifact hosted by PyPI.
-> These can be found in the "Download files" section of the project's page,
-> e.g: https://pypi.org/project/sigstore/#files
-
+> The package to verify can be passed either as a `pypi:` prefixed filename (e.g:
+> 'pypi:sampleproject-1.0.0-py3-none-any.whl'), or as a direct URL
+> to the artifact hosted by PyPI.
 ```bash
+pypi-attestations verify pypi --repository https://github.com/sigstore/sigstore-python \
+  pypi:sigstore-3.6.1-py3-none-any.whl
+
+# or alternatively:
 pypi-attestations verify pypi --repository https://github.com/sigstore/sigstore-python \
   https://files.pythonhosted.org/packages/70/f5/324edb6a802438e97e289992a41f81bb7a58a1cda2e49439e7e48896649e/sigstore-3.6.1-py3-none-any.whl
 ```
 
-This command downloads the artifact from the given URL and gets its provenance
-from PyPI. The artifact is then verified against the provenance, while also
-checking that the provenance's signing identity matches the repository specified
-by the user.
+This command downloads the artifact and its provenance from PyPI. The artifact 
+is then verified against the provenance, while also checking that the provenance's 
+signing identity matches the repository specified by the user.
 
 
 [PEP 740]: https://peps.python.org/pep-0740/
diff --git a/src/pypi_attestations/_cli.py b/src/pypi_attestations/_cli.py
@@ -37,6 +37,8 @@
     from collections.abc import Iterable
     from typing import NoReturn
 
+    from rfc3986 import URIReference
+
 logging.basicConfig(format="%(message)s", datefmt="[%X]", handlers=[logging.StreamHandler()])
 _logger = logging.getLogger(__name__)
 _logger.setLevel(logging.INFO)
@@ -134,10 +136,11 @@ def _parser() -> argparse.ArgumentParser:
     verify_pypi_command = verify_subcommands.add_parser(name="pypi", help="Verify a PyPI release")
 
     verify_pypi_command.add_argument(
-        "distribution_url",
-        metavar="URL_PYPI_FILE",
+        "distribution_file",
+        metavar="PYPI_FILE",
         type=str,
-        help='URL of the PyPI file to verify, i.e: "https://files.pythonhosted.org/..."',
+        help="PyPI file to verify, can be either: (1) pypi:$FILE_NAME (e.g. "
+        "pypi:sampleproject-1.0.0.tar.gz) or (2) A direct URL to files.pythonhosted.org",
     )
 
     verify_pypi_command.add_argument(
@@ -230,6 +233,61 @@ def _download_file(url: str, dest: Path) -> None:
             _die(f"Error downloading file: {e}")
 
 
+def _get_direct_url_from_arg(arg: str) -> URIReference:
+    """Parse the artifact argument for the `verify pypi` subcommand.
+
+    The argument can be:
+    - A pypi: prefixed filename (e.g. pypi:sampleproject-1.0.0.tar.gz)
+    - A direct URL to a PyPI-hosted artifact
+    """
+    direct_url = None
+
+    if arg.startswith("pypi:"):
+        file_name = arg[5:]
+        try:
+            if file_name.endswith(".tar.gz") or file_name.endswith(".zip"):
+                pkg_name, _ = parse_sdist_filename(file_name)
+            elif file_name.endswith(".whl"):
+                pkg_name, _, _, _ = parse_wheel_filename(file_name)
+            else:
+                _die("File should be a wheel (*.whl) or a source distribution (*.zip or *.tar.gz)")
+        except (InvalidSdistFilename, InvalidWheelFilename) as e:
+            _die(f"Invalid distribution filename: {e}")
+
+        provenance_url = f"https://pypi.org/simple/{pkg_name}"
+        response = requests.get(
+            provenance_url, headers={"Accept": "application/vnd.pypi.simple.v1+json"}
+        )
+        try:
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as e:
+            _die(f"Error trying to get information for '{pkg_name}' from PyPI: {e}")
+
+        response_json = response.json()
+        for file_json in response_json.get("files", []):
+            if file_json.get("filename", "") == file_name:
+                direct_url = file_json.get("url", "")
+                break
+        if not direct_url:
+            _die(f"Could not find the artifact '{file_name}' on PyPI")
+    else:
+        direct_url = arg
+
+    validator = (
+        validators.Validator()
+        .allow_schemes("https")
+        .allow_hosts("files.pythonhosted.org")
+        .require_presence_of("scheme", "host")
+    )
+    try:
+        pypi_url = uri_reference(direct_url)
+        validator.validate(pypi_url)
+    except exceptions.RFC3986Exception as e:
+        _die(f"Unsupported/invalid URL: {e}")
+
+    return pypi_url
+
+
 def _get_provenance_from_pypi(filename: str) -> Provenance:
     """Use PyPI's integrity API to get a distribution's provenance."""
     try:
@@ -425,17 +483,7 @@ def _verify_pypi(args: argparse.Namespace) -> None:
     the provenance file hosted on PyPI (if any), and against the repository URL
     passed by the user as a CLI argument.
     """
-    validator = (
-        validators.Validator()
-        .allow_schemes("https")
-        .allow_hosts("files.pythonhosted.org")
-        .require_presence_of("scheme", "host")
-    )
-    try:
-        pypi_url = uri_reference(args.distribution_url)
-        validator.validate(pypi_url)
-    except exceptions.RFC3986Exception as e:
-        _die(f"Unsupported/invalid URL: {e}")
+    pypi_url = _get_direct_url_from_arg(args.distribution_file)
 
     with TemporaryDirectory() as temp_dir:
         dist_filename = pypi_url.path.split("/")[-1]
diff --git a/test/test_cli.py b/test/test_cli.py
@@ -39,6 +39,8 @@
 pypi_sdist_url = "https://files.pythonhosted.org/packages/db/89/b982115aabe1068fd581d83d2a0b26b78e1e7ce6184e75003d173e15c0b3/sigstore-3.6.1.tar.gz"
 pypi_wheel_filename = pypi_wheel_url.split("/")[-1]
 pypi_sdist_filename = pypi_sdist_url.split("/")[-1]
+pypi_wheel_abbrev = f"sigstore/{pypi_wheel_filename}"
+pypi_sdist_abbrev = f"sigstore/{pypi_sdist_filename}"
 
 
 def run_main_with_command(cmd: list[str]) -> None:
@@ -366,35 +368,33 @@ def test_validate_files(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> Non
 
 
 @online
-def test_verify_pypi_command(caplog: pytest.LogCaptureFixture) -> None:
-    # Happy path wheel
-    run_main_with_command(
-        [
-            "verify",
-            "pypi",
-            "--repository",
-            "https://github.com/sigstore/sigstore-python",
-            pypi_wheel_url,
-        ]
-    )
-    assert f"OK: {pypi_wheel_filename}" in caplog.text
-
-    caplog.clear()
-
-    # Happy path sdist
+@pytest.mark.parametrize(
+    "url_argument, filename",
+    [
+        (pypi_wheel_url, pypi_wheel_filename),
+        (pypi_sdist_url, pypi_sdist_filename),
+        (f"pypi:{pypi_wheel_filename}", pypi_wheel_filename),
+        (f"pypi:{pypi_sdist_filename}", pypi_sdist_filename),
+    ],
+)
+def test_verify_pypi_command(
+    caplog: pytest.LogCaptureFixture, url_argument: str, filename: str
+) -> None:
+    # Happy path
     run_main_with_command(
         [
             "verify",
             "pypi",
             "--repository",
             "https://github.com/sigstore/sigstore-python",
-            pypi_sdist_url,
+            url_argument,
         ]
     )
-    assert f"OK: {pypi_sdist_filename}" in caplog.text
+    assert f"OK: {filename}" in caplog.text
 
-    caplog.clear()
 
+@online
+def test_verify_pypi_command_env_fail(caplog: pytest.LogCaptureFixture) -> None:
     with pytest.raises(SystemExit):
         # Failure from the Sigstore environment
         run_main_with_command(
@@ -468,7 +468,7 @@ def test_verify_pypi_invalid_url(
     assert "Unsupported/invalid URL" in caplog.text
 
 
-def test_verify_pypi_invalid_file_name(
+def test_verify_pypi_invalid_file_name_url(
     caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     # Failure because file is neither a wheer nor a sdist
@@ -503,6 +503,41 @@ def test_verify_pypi_invalid_file_name(
     assert "Invalid wheel filename" in caplog.text
 
 
+def test_verify_pypi_invalid_sdist_filename_pypi(
+    caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    # Failure because file is neither a wheer nor a sdist
+    monkeypatch.setattr(pypi_attestations._cli, "_download_file", lambda url, dest: None)
+    with pytest.raises(SystemExit):
+        run_main_with_command(
+            [
+                "verify",
+                "pypi",
+                "--repository",
+                "https://github.com/sigstore/sigstore-python",
+                f"pypi:{pypi_wheel_filename}.invalid_ext",
+            ]
+        )
+    assert (
+        "File should be a wheel (*.whl) or a source distribution (*.zip or *.tar.gz)" in caplog.text
+    )
+
+    caplog.clear()
+
+    """Test that invalid sdist filenames are properly handled."""
+    with pytest.raises(SystemExit):
+        run_main_with_command(
+            [
+                "verify",
+                "pypi",
+                "--repository",
+                "https://github.com/sigstore/sigstore-python",
+                "pypi:invalid-sdist-name.tar.gz",  # Invalid sdist filename format
+            ]
+        )
+    assert "Invalid distribution filename:" in caplog.text
+
+
 @online
 def test_verify_pypi_validation_fails(
     caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch
@@ -562,6 +597,44 @@ def test_verify_pypi_error_getting_provenance(
     assert expected_error in caplog.text
 
 
+def test_verify_pypi_error_finding_package_info(
+    caplog: pytest.LogCaptureFixture,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    response = stub(raise_for_status=raiser(requests.HTTPError("myerror")))
+    monkeypatch.setattr(requests, "get", lambda url, headers: response)
+    with pytest.raises(SystemExit):
+        run_main_with_command(
+            [
+                "verify",
+                "pypi",
+                "--repository",
+                "https://github.com/sigstore/sigstore-python",
+                "pypi:somefile-1.0.0.tar.gz",
+            ]
+        )
+    assert "Error trying to get information for 'somefile' from PyPI: myerror" in caplog.text
+
+
+def test_verify_pypi_error_finding_artifact_url(
+    caplog: pytest.LogCaptureFixture,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    response = stub(raise_for_status=lambda: None, json=lambda: {"files": []})
+    monkeypatch.setattr(requests, "get", lambda url, headers: response)
+    with pytest.raises(SystemExit):
+        run_main_with_command(
+            [
+                "verify",
+                "pypi",
+                "--repository",
+                "https://github.com/sigstore/sigstore-python",
+                "pypi:somefile-1.0.0.tar.gz",
+            ]
+        )
+    assert "Could not find the artifact 'somefile-1.0.0.tar.gz' on PyPI" in caplog.text
+
+
 def test_verify_pypi_error_validating_provenance(
     caplog: pytest.LogCaptureFixture,
     monkeypatch: pytest.MonkeyPatch,
