fix: poetry error on source and parsing pyproject.toml

Now we auto-add the Safety source and also we fixed the issues parsing Poetry projects.

Author: yeisonvargasf

safety/tool/base.py

@@ -344,7 +344,7 @@ def can_handle(self, args: List[str]) -> bool:
         """
         Check if this parser can handle the given arguments
         """
-        if not args or len(args) < 2:
+        if not args or len(args) < 1:
             return False
 
         return args[0].lower() in self.intention_mapping

safety/tool/poetry/command.py

@@ -1,10 +1,12 @@
 from pathlib import Path
 import sys
 from typing import TYPE_CHECKING, List, Optional, Tuple
-
+import logging
 import typer
 
-from .constants import MSG_MISSING_SAFETY_SOURCE
+from safety.tool.utils import PoetryPyprojectConfigurator
+
+from .constants import MSG_SAFETY_SOURCE_ADDED, MSG_SAFETY_SOURCE_NOT_ADDED
 from .parser import PoetryParser
 
 from ..base import BaseCommand, ToolIntentionType
@@ -24,6 +26,8 @@
 else:
     import tomli as tomllib
 
+logger = logging.getLogger(__name__)
+
 
 class PoetryCommand(BaseCommand):
     """
@@ -74,23 +78,7 @@ def get_package_list_command(self) -> List[str]:
         """
         return ["poetry", "run", "pip", "list", "--format=json"]
 
-    @classmethod
-    def from_args(cls, args: List[str], **kwargs):
-        parser = PoetryParser()
-
-        if intention := parser.parse(args):
-            if intention.intention_type is ToolIntentionType.ADD_PACKAGE:
-                return PoetryAddCommand(args, intention=intention, **kwargs)
-
-        return PoetryGenericCommand(args, **kwargs)
-
-
-class PoetryGenericCommand(PoetryCommand):
-    pass
-
-
-class PoetryAddCommand(PoetryCommand):
-    def has_safety_source_in_pyproject(self) -> bool:
+    def _has_safety_source_in_pyproject(self) -> bool:
         """
         Check if 'safety' source exists in pyproject.toml
         """
@@ -113,6 +101,55 @@ def has_safety_source_in_pyproject(self) -> bool:
         except (FileNotFoundError, KeyError, tomllib.TOMLDecodeError):
             return False
 
+    def before(self, ctx: typer.Context):
+        super().before(ctx)
+
+        if self._intention and self._intention.intention_type in [
+            ToolIntentionType.SYNC_PACKAGES,
+            ToolIntentionType.ADD_PACKAGE,
+        ]:
+            if not self._has_safety_source_in_pyproject():
+                org_slug = None
+                try:
+                    data = ctx.obj.auth.client.initialize()
+                    org_slug = data.get("organization-data", {}).get("slug")
+                except Exception:
+                    logger.exception(
+                        "Unable to pull the org slug from the initialize endpoint."
+                    )
+
+                try:
+                    configurator = PoetryPyprojectConfigurator()
+                    prj_slug = ctx.obj.project.id if ctx.obj.project else None
+                    if configurator.configure(
+                        Path("pyproject.toml"), org_slug, prj_slug
+                    ):
+                        console.print(
+                            MSG_SAFETY_SOURCE_ADDED,
+                        )
+                except Exception:
+                    logger.exception("Unable to configure the pyproject.toml file.")
+                    console.print(
+                        MSG_SAFETY_SOURCE_NOT_ADDED,
+                    )
+
+    @classmethod
+    def from_args(cls, args: List[str], **kwargs):
+        parser = PoetryParser()
+
+        if intention := parser.parse(args):
+            kwargs["intention"] = intention
+            if intention.intention_type is ToolIntentionType.ADD_PACKAGE:
+                return PoetryAddCommand(args, **kwargs)
+
+        return PoetryGenericCommand(args, **kwargs)
+
+
+class PoetryGenericCommand(PoetryCommand):
+    pass
+
+
+class PoetryAddCommand(PoetryCommand):
     def patch_source_option(
         self, args: List[str], new_source: str = "safety"
     ) -> Tuple[Optional[str], List[str]]:
@@ -146,11 +183,5 @@ def patch_source_option(
     def before(self, ctx: typer.Context):
         super().before(ctx)
 
-        if not self.has_safety_source_in_pyproject():
-            console.print(
-                MSG_MISSING_SAFETY_SOURCE,
-            )
-            sys.exit(1)
-
         _, modified_args = self.patch_source_option(self._args)
         self._args = modified_args

safety/tool/poetry/constants.py

@@ -1 +1,4 @@
-MSG_MISSING_SAFETY_SOURCE = "\nError: Safety Firewall source is not configured in pyproject.toml. Continuing with installation would result in you not being protected from malicious or insecure packages. Run 'safety init' to fix this."
+MSG_SAFETY_SOURCE_NOT_ADDED = "\nError: Safety Firewall could not be added as a source in your pyproject.toml file. You will not be protected from malicious or insecure packages. Please run `safety init` to fix this."
+MSG_SAFETY_SOURCE_ADDED = (
+    "\nSafety Firewall has been added as a source to protect this codebase"
+)

safety/tool/poetry/main.py

@@ -38,10 +38,17 @@ def is_installed(cls) -> bool:
     def is_poetry_project_file(cls, file: Path) -> bool:
         try:
             cfg = tomllib.loads(file.read_text())
-            return cfg.get("build-system", {}).get("requires") in [
-                ["poetry-core"],
-                "poetry-core",
-            ]
+
+            # First check: tool.poetry section (most definitive)
+            if "tool" in cfg and "poetry" in cfg.get("tool", {}):
+                return True
+
+            # Extra check on build-system section
+            build_backend = cfg.get("build-system", {}).get("build-backend", "")
+            if build_backend and "poetry.core" in build_backend:
+                return True
+
+            return False
         except (IOError, ValueError):
             return False
 

safety/tool/poetry/parser.py

@@ -20,6 +20,7 @@ def intention_mapping(self) -> Dict[str, ToolIntentionType]:
             "show": ToolIntentionType.LIST_PACKAGES,
             "init": ToolIntentionType.INIT_PROJECT,
             "build": ToolIntentionType.BUILD_PROJECT,
+            "install": ToolIntentionType.SYNC_PACKAGES,
         }
 
     def parse(self, args: List[str]) -> Optional[CommandToolIntention]: