Support pyproject.toml

CHANGELOG.md

@@ -14,6 +14,10 @@ All versions prior to 0.0.9 are untracked.
   packages or parsed requirements (via `-r`) that are marked as editable
   ([#244](https://github.com/trailofbits/pip-audit/pull/244))
 
+* CLI: `pip-audit` can audit projects that list their dependencies in
+  `pyproject.toml` files, via `pip-audit <dir>`
+  ([#246](https://github.com/trailofbits/pip-audit/pull/246))
+
 ## [2.0.0] - 2022-02-18
 
 ### Added

README.md

@@ -74,9 +74,14 @@ usage: pip-audit [-h] [-V] [-l] [-r REQUIREMENTS] [-f FORMAT] [-s SERVICE]
                  [--path PATHS] [-v] [--fix] [--require-hashes]
                  [--index-url INDEX_URL] [--extra-index-url EXTRA_INDEX_URLS]
                  [--skip-editable]
+                 [project_path]
 
 audit the Python environment for dependencies with known vulnerabilities
 
+positional arguments:
+  project_path          audit a local Python project at the given path
+                        (default: None)
+
 optional arguments:
   -h, --help            show this help message and exit
   -V, --version         show program's version number and exit
@@ -175,6 +180,13 @@ $ pip-audit -r ./requirements.txt -l
 No known vulnerabilities found
 ```
 
+Audit dependencies for a local Python project:
+```
+$ pip-audit .
+No known vulnerabilities found
+```
+`pip-audit` searches the provided path for various Python "project" files. At the moment, only `pyproject.toml` is supported.
+
 Audit dependencies when there are vulnerabilities present:
 ```
 $ pip-audit

pip_audit/_cli.py

@@ -17,6 +17,7 @@
     PYPI_URL,
     DependencySource,
     PipSource,
+    PyProjectSource,
     RequirementSource,
     ResolveLibResolver,
 )
@@ -159,6 +160,9 @@ def _parser() -> argparse.ArgumentParser:
         dest="requirements",
         help="audit the given requirements file; this option can be used multiple times",
     )
+    dep_source_args.add_argument(
+        "project_path", type=Path, nargs="?", help="audit a local Python project at the given path"
+    )
     parser.add_argument(
         "-f",
         "--format",
@@ -274,6 +278,17 @@ def _parse_args(parser: argparse.ArgumentParser) -> argparse.Namespace:
     return parser.parse_args()
 
 
+def _dep_source_from_project_path(project_path: Path, state: AuditState) -> DependencySource:
+    # Check for a `pyproject.toml`
+    pyproject_path = project_path / "pyproject.toml"
+    if pyproject_path.is_file():
+        return PyProjectSource(pyproject_path, ResolveLibResolver(), state)
+
+    # TODO: Checks for setup.py and other project files will go here.
+
+    _fatal(f"couldn't find a supported project file in {project_path}")
+
+
 def audit() -> None:
     """
     The primary entrypoint for `pip-audit`.
@@ -306,8 +321,8 @@ def audit() -> None:
         state = stack.enter_context(AuditState(members=actors))
 
         source: DependencySource
+        index_urls = [args.index_url] + args.extra_index_urls
         if args.requirements is not None:
-            index_urls = [args.index_url] + args.extra_index_urls
             req_files: List[Path] = [Path(req.name) for req in args.requirements]
             # TODO: This is a leaky abstraction; we should construct the ResolveLibResolver
             # within the RequirementSource instead of in-line here.
@@ -319,6 +334,12 @@ def audit() -> None:
                 require_hashes=args.require_hashes,
                 state=state,
             )
+        elif args.project_path is not None:
+            # NOTE: We'll probably want to support --skip-editable here,
+            # once PEP 660 is more widely supported: https://www.python.org/dev/peps/pep-0660/
+
+            # Determine which kind of project file exists in the project path
+            source = _dep_source_from_project_path(args.project_path, state)
         else:
             source = PipSource(
                 local=args.local, paths=args.paths, skip_editable=args.skip_editable, state=state

pip_audit/_dependency_source/__init__.py

@@ -10,6 +10,7 @@
     DependencySourceError,
 )
 from .pip import PipSource, PipSourceError
+from .pyproject import PyProjectSource
 from .requirement import RequirementSource
 from .resolvelib import PYPI_URL, ResolveLibResolver
 
@@ -22,6 +23,7 @@
     "DependencySourceError",
     "PipSource",
     "PipSourceError",
+    "PyProjectSource",
     "RequirementSource",
     "ResolveLibResolver",
 ]

pip_audit/_dependency_source/pyproject.py

@@ -0,0 +1,147 @@
+"""
+Collect dependencies from `pyproject.toml` files.
+"""
+
+import logging
+import os
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Iterator, List, Set, cast
+
+import toml
+from packaging.requirements import Requirement
+from packaging.specifiers import SpecifierSet
+
+from pip_audit._dependency_source import (
+    DependencyFixError,
+    DependencyResolver,
+    DependencyResolverError,
+    DependencySource,
+    DependencySourceError,
+)
+from pip_audit._fix import ResolvedFixVersion
+from pip_audit._service import Dependency, ResolvedDependency, SkippedDependency
+from pip_audit._state import AuditState
+
+logger = logging.getLogger(__name__)
+
+
+class PyProjectSource(DependencySource):
+    """
+    Wraps `pyproject.toml` dependency resolution as a dependency source.
+    """
+
+    def __init__(
+        self, filename: Path, resolver: DependencyResolver, state: AuditState = AuditState()
+    ) -> None:
+        """
+        Create a new `PyProjectSource`.
+
+        `filename` provides a path to a `pyproject.toml` file
+
+        `resolver` is the `DependencyResolver` to use.
+
+        `state` is an `AuditState` to use for state callbacks.
+        """
+        self.filename = filename
+        self.resolver = resolver
+        self.state = state
+
+    def collect(self) -> Iterator[Dependency]:
+        """
+        Collect all of the dependencies discovered by this `PyProjectSource`.
+
+        Raises a `PyProjectSourceError` on any errors.
+        """
+
+        collected: Set[Dependency] = set()
+        with self.filename.open("r") as f:
+            pyproject_data = toml.load(f)
+
+            project = pyproject_data.get("project")
+            if project is None:
+                raise PyProjectSourceError(
+                    f"pyproject file {self.filename} does not contain `project` section"
+                )
+
+            deps = project.get("dependencies")
+            if deps is None:
+                # Projects without dependencies aren't an error case
+                logger.warning(
+                    f"pyproject file {self.filename} does not contain `dependencies` list"
+                )
+                return
+
+            reqs: List[Requirement] = [Requirement(dep) for dep in deps]
+            try:
+                for _, deps in self.resolver.resolve_all(iter(reqs)):
+                    for dep in deps:
+                        # Don't allow duplicate dependencies to be returned
+                        if dep in collected:
+                            continue
+
+                        if dep.is_skipped():  # pragma: no cover
+                            dep = cast(SkippedDependency, dep)
+                            self.state.update_state(f"Skipping {dep.name}: {dep.skip_reason}")
+                        else:
+                            dep = cast(ResolvedDependency, dep)
+                            self.state.update_state(f"Collecting {dep.name} ({dep.version})")
+
+                        collected.add(dep)
+                        yield dep
+            except DependencyResolverError as dre:
+                raise PyProjectSourceError("dependency resolver raised an error") from dre
+
+    def fix(self, fix_version: ResolvedFixVersion) -> None:
+        """
+        Fixes a dependency version for this `PyProjectSource`.
+        """
+
+        with self.filename.open("r+") as f, NamedTemporaryFile(mode="r+", delete=False) as tmp:
+            pyproject_data = toml.load(f)
+
+            project = pyproject_data.get("project")
+            if project is None:
+                raise PyProjectFixError(
+                    f"pyproject file {self.filename} does not contain `project` section"
+                )
+
+            deps = project.get("dependencies")
+            if deps is None:
+                # Projects without dependencies aren't an error case
+                logger.warning(
+                    f"pyproject file {self.filename} does not contain `dependencies` list"
+                )
+                return
+
+            reqs: List[Requirement] = [Requirement(dep) for dep in deps]
+            for i in range(len(reqs)):
+                # When we find a requirement that matches the provided fix version, we need to edit
+                # the requirement's specifier and then write it back to the underlying TOML data.
+                req = reqs[i]
+                if (
+                    req.name == fix_version.dep.name
+                    and req.specifier.contains(fix_version.dep.version)
+                    and not req.specifier.contains(fix_version.version)
+                ):
+                    req.specifier = SpecifierSet(f"=={fix_version.version}")
+                    deps[i] = str(req)
+                assert req.marker is None or req.marker.evaluate()
+
+            # Now dump the new edited TOML to the temporary file.
+            toml.dump(pyproject_data, tmp)
+
+            # And replace the original `pyproject.toml` file.
+            os.replace(tmp.name, self.filename)
+
+
+class PyProjectSourceError(DependencySourceError):
+    """A `pyproject.toml` specific `DependencySourceError`."""
+
+    pass
+
+
+class PyProjectFixError(DependencyFixError):
+    """A `pyproject.toml` specific `DependencyFixError`."""
+
+    pass