Adding --osv-url argument to allow use of private OSV vulnerability services (#810)

* Adding --osv-url argument to allow use of private OSV vulnerability
services

#805

* cleanup, refactoring, add CHANGELOG

---------

Co-authored-by: William Woodruff <[email protected]>
Co-authored-by: William Woodruff <[email protected]>

Author: davidjmemmett

.gitignore

@@ -6,3 +6,4 @@ __pycache__/
 html/
 dist/
 .python-version
+/.pytest_cache/

CHANGELOG.md

@@ -8,8 +8,18 @@ All versions prior to 0.0.9 are untracked.
 
 ## [Unreleased]
 
+### Added
+
+* `pip-audit` now supports the `--osv-url URL` flag, which can be used to
+  retrieve vulnerabilities from a custom OSV service. This is useful for
+  organizations that host their own mirror of the OSV database, or that
+  have custom OSV records
+  ([#810](https://github.com/pypa/pip-audit/pull/810))
+
 ## [2.9.0]
 
+### Added
+
 * `pip-audit` now supports [PEP 751](https://peps.python.org/pep-0751/)
   lockfiles. These lockfiles can be audited in "project" mode by
   passing `--locked` to `pip-audit`

README.md

@@ -132,13 +132,13 @@ python -m pip_audit --help
 <!-- @begin-pip-audit-help@ -->
 ```
 usage: pip-audit [-h] [-V] [-l] [-r REQUIREMENT] [--locked] [-f FORMAT]
-                 [-s SERVICE] [-d] [-S] [--desc [{on,off,auto}]]
-                 [--aliases [{on,off,auto}]] [--cache-dir CACHE_DIR]
-                 [--progress-spinner {on,off}] [--timeout TIMEOUT]
-                 [--path PATH] [-v] [--fix] [--require-hashes]
-                 [--index-url INDEX_URL] [--extra-index-url URL]
-                 [--skip-editable] [--no-deps] [-o FILE] [--ignore-vuln ID]
-                 [--disable-pip]
+                 [-s SERVICE] [--osv-url OSV_URL] [-d] [-S]
+                 [--desc [{on,off,auto}]] [--aliases [{on,off,auto}]]
+                 [--cache-dir CACHE_DIR] [--progress-spinner {on,off}]
+                 [--timeout TIMEOUT] [--path PATH] [-v] [--fix]
+                 [--require-hashes] [--index-url INDEX_URL]
+                 [--extra-index-url URL] [--skip-editable] [--no-deps]
+                 [-o FILE] [--ignore-vuln ID] [--disable-pip]
                  [project_path]
 
 audit the Python environment for dependencies with known vulnerabilities
@@ -165,6 +165,8 @@ optional arguments:
   -s SERVICE, --vulnerability-service SERVICE
                         the vulnerability service to audit dependencies
                         against (choices: osv, pypi) (default: pypi)
+  --osv-url OSV_URL     URL to use for the OSV API instead of the default
+                        (default: https://api.osv.dev/v1/query)
   -d, --dry-run         without `--fix`: collect all dependencies but do not
                         perform the auditing step; with `--fix`: perform the
                         auditing step but do not perform any fixes (default:

pip_audit/_cli.py

@@ -32,9 +32,9 @@
     MarkdownFormat,
     VulnerabilityFormat,
 )
-from pip_audit._service import OsvService, PyPIService, VulnerabilityService
+from pip_audit._service import OsvService, PyPIService
 from pip_audit._service.interface import ConnectionError as VulnServiceConnectionError
-from pip_audit._service.interface import ResolvedDependency, SkippedDependency
+from pip_audit._service.interface import ResolvedDependency, SkippedDependency, VulnerabilityService
 from pip_audit._state import AuditSpinner, AuditState
 from pip_audit._util import assert_never
 
@@ -100,14 +100,6 @@ class VulnerabilityServiceChoice(str, enum.Enum):
     Osv = "osv"
     Pypi = "pypi"
 
-    def to_service(self, timeout: int, cache_dir: Path | None) -> VulnerabilityService:
-        if self is VulnerabilityServiceChoice.Osv:
-            return OsvService(cache_dir, timeout)
-        elif self is VulnerabilityServiceChoice.Pypi:
-            return PyPIService(cache_dir, timeout)
-        else:
-            assert_never(self)  # pragma: no cover
-
     def __str__(self) -> str:
         return self.value
 
@@ -249,6 +241,14 @@ def _parser() -> argparse.ArgumentParser:  # pragma: no cover
             VulnerabilityServiceChoice,
         ),
     )
+    parser.add_argument(
+        "--osv-url",
+        type=str,
+        metavar="OSV_URL",
+        dest="osv_url",
+        default=os.environ.get("PIP_AUDIT_OSV_URL", OsvService.DEFAULT_OSV_URL),
+        help="URL to use for the OSV API instead of the default",
+    )
     parser.add_argument(
         "-d",
         "--dry-run",
@@ -438,7 +438,14 @@ def audit() -> None:  # pragma: no cover
     parser = _parser()
     args = _parse_args(parser)
 
-    service = args.vulnerability_service.to_service(args.timeout, args.cache_dir)
+    service: VulnerabilityService
+    if args.vulnerability_service is VulnerabilityServiceChoice.Osv:
+        service = OsvService(cache_dir=args.cache_dir, timeout=args.timeout, osv_url=args.osv_url)
+    elif args.vulnerability_service is VulnerabilityServiceChoice.Pypi:
+        service = PyPIService(cache_dir=args.cache_dir, timeout=args.timeout)
+    else:
+        assert_never(args.vulnerability_service)  # pragma: no cover
+
     output_desc = args.desc.to_bool(args.format)
     output_aliases = args.aliases.to_bool(args.format)
     formatter = args.format.to_format(output_desc, output_aliases)

pip_audit/_service/osv.py

@@ -31,7 +31,14 @@ class OsvService(VulnerabilityService):
     package vulnerability information.
     """
 
-    def __init__(self, cache_dir: Path | None = None, timeout: int | None = None):
+    DEFAULT_OSV_URL = "https://api.osv.dev/v1/query"
+
+    def __init__(
+        self,
+        cache_dir: Path | None = None,
+        timeout: int | None = None,
+        osv_url: str = DEFAULT_OSV_URL,
+    ):
         """
         Create a new `OsvService`.
 
@@ -43,6 +50,7 @@ def __init__(self, cache_dir: Path | None = None, timeout: int | None = None):
         """
         self.session = caching_session(cache_dir, use_pip=False)
         self.timeout = timeout
+        self.osv_url = osv_url
 
     def query(self, spec: Dependency) -> tuple[Dependency, list[VulnerabilityResult]]:
         """
@@ -54,14 +62,13 @@ def query(self, spec: Dependency) -> tuple[Dependency, list[VulnerabilityResult]
             return spec, []
         spec = cast(ResolvedDependency, spec)
 
-        url = "https://api.osv.dev/v1/query"
         query = {
             "package": {"name": spec.canonical_name, "ecosystem": "PyPI"},
             "version": str(spec.version),
         }
         try:
             response: requests.Response = self.session.post(
-                url=url,
+                url=self.osv_url,
                 data=json.dumps(query),
                 timeout=self.timeout,
             )

pip_audit/_service/pypi.py

@@ -32,7 +32,9 @@ class PyPIService(VulnerabilityService):
     package vulnerability information.
     """
 
-    def __init__(self, cache_dir: Path | None = None, timeout: int | None = None) -> None:
+    def __init__(
+        self, cache_dir: Path | None = None, timeout: int | None = None, **kwargs: dict
+    ) -> None:
         """
         Create a new `PyPIService`.
 

test/test_cli.py

@@ -27,10 +27,6 @@ def test_str(self):
 
 
 class TestVulnerabilityServiceChoice:
-    def test_to_service_is_exhaustive(self, cache_dir):
-        for choice in VulnerabilityServiceChoice:
-            assert choice.to_service(0, cache_dir) is not None
-
     def test_str(self):
         for choice in VulnerabilityServiceChoice:
             assert str(choice) == choice.value