Merge pull request #322 from vipyrsec/check-pypi-html

Allow reporting on quarantined packages

src/mainframe/dependencies.py

@@ -18,6 +18,10 @@ def get_pypi_client() -> PyPIServices:
     return PyPIServices(http_client)
 
 
+def get_httpx_client(request: Request) -> httpx.Client:
+    return request.app.state.http_session
+
+
 def get_rules(request: Request) -> Rules:
     return request.app.state.rules
 

src/mainframe/endpoints/report.py

@@ -5,13 +5,12 @@
 import structlog
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.encoders import jsonable_encoder
-from letsbuilda.pypi import PackageNotFoundError, PyPIServices
 from sqlalchemy import select
 from sqlalchemy.orm import Session, joinedload
 
 from mainframe.constants import mainframe_settings
 from mainframe.database import get_db
-from mainframe.dependencies import get_pypi_client, validate_token
+from mainframe.dependencies import get_httpx_client, validate_token
 from mainframe.json_web_token import AuthenticationData
 from mainframe.models.orm import Scan
 from mainframe.models.schemas import (
@@ -138,12 +137,11 @@ def _validate_additional_information(body: ReportPackageBody, scan: Scan):
             raise error
 
 
-def _validate_pypi(name: str, version: str, pypi_client: PyPIServices):
+def _validate_pypi(name: str, version: str, http_client: httpx.Client):
     log = logger.bind(package={"name": name, "version": version})
 
-    try:
-        pypi_client.get_package_metadata(name, version)
-    except PackageNotFoundError:
+    response = http_client.get(f"https://pypi.org/project/{name}")
+    if response.status_code == 404:
         error = HTTPException(404, detail="Package not found on PyPI")
         log.error("Package not found on PyPI", tag="package_not_found_pypi")
         raise error
@@ -160,7 +158,7 @@ def report_package(
     body: ReportPackageBody,
     session: Annotated[Session, Depends(get_db)],
     auth: Annotated[AuthenticationData, Depends(validate_token)],
-    pypi_client: Annotated[PyPIServices, Depends(get_pypi_client)],
+    httpx_client: Annotated[httpx.Client, Depends(get_httpx_client)],
 ):
     """
     Report a package to PyPI.
@@ -206,7 +204,7 @@ def report_package(
 
     # If execution reaches here, we must have found a matching scan in our
     # database. Check if the package we want to report exists on PyPI.
-    _validate_pypi(name, version, pypi_client)
+    _validate_pypi(name, version, httpx_client)
 
     rules_matched: list[str] = [rule.name for rule in scan.rules]
 
@@ -220,7 +218,7 @@ def report_package(
             additional_information=body.additional_information,
         )
 
-        httpx.post(f"{mainframe_settings.reporter_url}/report/email", json=jsonable_encoder(report))
+        httpx_client.post(f"{mainframe_settings.reporter_url}/report/email", json=jsonable_encoder(report))
     else:
         # We previously checked this condition, but the typechecker isn't smart
         # enough to figure that out
@@ -233,7 +231,7 @@ def report_package(
             extra=dict(yara_rules=rules_matched),
         )
 
-        httpx.post(f"{mainframe_settings.reporter_url}/report/{name}", json=jsonable_encoder(report))
+        httpx_client.post(f"{mainframe_settings.reporter_url}/report/{name}", json=jsonable_encoder(report))
 
     with session.begin():
         scan.reported_by = auth.subject

tests/test_report.py

@@ -7,9 +7,6 @@
 import pytest
 from fastapi import HTTPException
 from fastapi.encoders import jsonable_encoder
-from letsbuilda.pypi import PyPIServices
-from letsbuilda.pypi.exceptions import PackageNotFoundError
-from pytest import MonkeyPatch
 from sqlalchemy import select
 from sqlalchemy.orm import Session, sessionmaker
 
@@ -79,7 +76,6 @@ def test_report(
     sm: sessionmaker[Session],
     db_session: Session,
     auth: AuthenticationData,
-    pypi_client: PyPIServices,
     body: ReportPackageBody,
     url: str,
     expected: EmailReport | ObservationReport,
@@ -107,11 +103,11 @@ def test_report(
     with db_session.begin():
         db_session.add(scan)
 
-    httpx.post = MagicMock()
+    mock_httpx_client = MagicMock()
 
-    report_package(body, sm(), auth, pypi_client)
+    report_package(body, sm(), auth, mock_httpx_client)
 
-    httpx.post.assert_called_once_with(url, json=jsonable_encoder(expected))
+    mock_httpx_client.post.assert_called_once_with(url, json=jsonable_encoder(expected))
 
     with sm() as sess, sess.begin():
         s = sess.scalar(select(Scan).where(Scan.name == "c").where(Scan.version == "1.0.0"))
@@ -121,18 +117,12 @@ def test_report(
     assert s.reported_at is not None
 
 
-def test_report_package_not_on_pypi(
-    pypi_client: PyPIServices,
-    monkeypatch: MonkeyPatch,
-):
-    # Make get_package_metadata always throw PackageNotFoundError to simulate an invalid package
-    def _side_effect(name: str, version: str):
-        raise PackageNotFoundError(name, version)
-
-    monkeypatch.setattr(pypi_client, "get_package_metadata", _side_effect)
+def test_report_package_not_on_pypi():
+    mock_httpx_client = MagicMock(spec=httpx.Client)
+    mock_httpx_client.configure_mock(**{"get.return_value.status_code": 404})
 
     with pytest.raises(HTTPException) as e:
-        _validate_pypi("c", "1.0.0", pypi_client)
+        _validate_pypi("c", "1.0.0", mock_httpx_client)
     assert e.value.status_code == 404