cli: Support friendlier syntax for verify pypi command

Signed-off-by: Facundo Tuesca <[email protected]>

Author: facutuesca

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,12 @@ 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 formatted as $PKG_NAME/$FILE_NAME, e.g: "
+        "sampleproject/sampleproject-1.0.0.tar.gz. Direct URLs to the hosted file are also "
+        "supported.",
     )
 
     verify_pypi_command.add_argument(
@@ -230,6 +234,52 @@ 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 either a direct URL to a PyPI-hosted artifact,
+    or a friendly-formatted alternative: '$PKG_NAME/$FILE_NAME'.
+    """
+    direct_url = None
+    components = arg.split("/")
+    # We support passing the file argument as $PKG_NAME/$FILE_NAME
+    if len(components) == 2:
+        pkg_name, file_name = components
+
+        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}' for '{pkg_name}'")
+    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 +475,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]

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),
+        (pypi_wheel_abbrev, pypi_wheel_filename),
+        (pypi_sdist_abbrev, 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(
@@ -562,6 +562,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",
+                "somepkg/somefile",
+            ]
+        )
+    assert "Error trying to get information for 'somepkg' 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",
+                "somepkg/somefile",
+            ]
+        )
+    assert "Could not find the artifact 'somefile' for 'somepkg'" in caplog.text
+
+
 def test_verify_pypi_error_validating_provenance(
     caplog: pytest.LogCaptureFixture,
     monkeypatch: pytest.MonkeyPatch,