feat: configurable pypi-reference

Author: Sergio Castillo

README.md

@@ -131,9 +131,15 @@ twyn run --config <file>
 All the configurations available through the command line are also supported in the config file.
 
 ```toml
-  [tool.twyn]
-  dependency_file="/my/path/requirements.txt"
-  selector_method="first_letter"
-  logging_level="debug"
-  allowlist=["my_package"]
+[tool.twyn]
+dependency_file="/my/path/requirements.txt"
+selector_method="first_letter"
+logging_level="debug"
+allowlist=["my_package"]
+pypi_reference="https://mirror-with-trusted-dependencies.com/file.json"
 ```
+
+> [!WARNING]
+> `twyn` will have a default reference URL for every source of trusted packages that is configurable.
+> If you want to protect yourself against spoofing attacks, it is recommended to set your own
+> reference url.

src/twyn/base/constants.py

@@ -9,6 +9,7 @@
 if TYPE_CHECKING:
     from twyn.dependency_parser.abstract_parser import AbstractParser
 
+
 SELECTOR_METHOD_MAPPING: dict[str, type[selectors.AbstractSelector]] = {
     "first-letter": selectors.FirstLetterExact,
     "nearby-letter": selectors.FirstLetterNearbyInKeyboard,
@@ -23,6 +24,7 @@
 DEFAULT_SELECTOR_METHOD = "all"
 DEFAULT_DEPENDENCY_FILE = "requirements.txt"
 DEFAULT_PROJECT_TOML_FILE = "pyproject.toml"
+DEFAULT_TOP_PYPI_PACKAGES = "https://hugovk.github.io/top-pypi-packages/top-pypi-packages-30-days.min.json"
 
 
 class AvailableLoggingLevels(enum.Enum):

src/twyn/core/config_handler.py

@@ -10,6 +10,7 @@
 from twyn.base.constants import (
     DEFAULT_PROJECT_TOML_FILE,
     DEFAULT_SELECTOR_METHOD,
+    DEFAULT_TOP_PYPI_PACKAGES,
     AvailableLoggingLevels,
 )
 from twyn.core.exceptions import (
@@ -29,6 +30,7 @@ class TwynConfiguration:
     selector_method: str
     logging_level: AvailableLoggingLevels
     allowlist: set[str]
+    pypi_reference: str
 
 
 @dataclass(frozen=True)
@@ -39,6 +41,7 @@ class ReadTwynConfiguration:
     selector_method: Optional[str]
     logging_level: Optional[AvailableLoggingLevels]
     allowlist: set[str]
+    pypi_reference: Optional[str]
 
 
 class ConfigHandler:
@@ -71,6 +74,7 @@ def resolve_config(
             selector_method=selector_method or twyn_config_data.get("selector_method", DEFAULT_SELECTOR_METHOD),
             logging_level=_get_logging_level(verbosity, twyn_config_data.get("logging_level")),
             allowlist=set(twyn_config_data.get("allowlist", set())),
+            pypi_reference=twyn_config_data.get("pypi_reference", DEFAULT_TOP_PYPI_PACKAGES),
         )
 
     def add_package_to_allowlist(self, package_name: str) -> None:
@@ -85,6 +89,7 @@ def add_package_to_allowlist(self, package_name: str) -> None:
             selector_method=config.selector_method,
             logging_level=config.logging_level,
             allowlist=config.allowlist | {package_name},
+            pypi_reference=config.pypi_reference,
         )
         self._write_config(toml, new_config)
         logger.info(f"Package '{package_name}' successfully added to allowlist")
@@ -101,6 +106,7 @@ def remove_package_from_allowlist(self, package_name: str) -> None:
             selector_method=config.selector_method,
             logging_level=config.logging_level,
             allowlist=config.allowlist - {package_name},
+            pypi_reference=config.pypi_reference,
         )
         self._write_config(toml, new_config)
         logger.info(f"Package '{package_name}' successfully removed from allowlist")
@@ -113,6 +119,7 @@ def _get_read_config(self, toml: TOMLDocument) -> ReadTwynConfiguration:
             selector_method=twyn_config_data.get("selector_method"),
             logging_level=twyn_config_data.get("logging_level"),
             allowlist=set(twyn_config_data.get("allowlist", set())),
+            pypi_reference=twyn_config_data.get("pypi_reference"),
         )
 
     def _write_config(self, toml: TOMLDocument, config: ReadTwynConfiguration) -> None:

src/twyn/main.py

@@ -41,7 +41,7 @@ def check_dependencies(
     _set_logging_level(config.logging_level)
 
     trusted_packages = TrustedPackages(
-        names=TopPyPiReference().get_packages(),
+        names=TopPyPiReference(source=config.pypi_reference).get_packages(),
         algorithm=EditDistance(),
         selector=get_candidate_selector(config.selector_method),
         threshold_class=SimilarityThreshold,

src/twyn/trusted_packages/constants.py

@@ -1,11 +1,3 @@
-from typing import NewType
-
-Url = NewType("Url", str)
-
-TOP_PYPI_PACKAGES = Url(
-    "https://hugovk.github.io/top-pypi-packages/top-pypi-packages-30-days.min.json"
-)
-
 
 ADJACENCY_MATRIX = {
     "1": ["2", "q", "w"],

src/twyn/trusted_packages/references.py

@@ -4,7 +4,6 @@
 
 import requests
 
-from twyn.trusted_packages.constants import TOP_PYPI_PACKAGES, Url
 from twyn.trusted_packages.exceptions import (
     EmptyPackagesListError,
     InvalidJSONError,
@@ -17,8 +16,8 @@
 class AbstractPackageReference(ABC):
     """Represents a reference to retrieve trusted packages from."""
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(self, source: str) -> None:
+        self.source = source
 
     @abstractmethod
     def get_packages(self) -> set[str]:
@@ -28,12 +27,6 @@ def get_packages(self) -> set[str]:
 class TopPyPiReference(AbstractPackageReference):
     """Top PyPi packages retrieved from an online source."""
 
-    def __init__(
-        self, source: Url = TOP_PYPI_PACKAGES, *args: Any, **kwargs: Any
-    ) -> None:
-        super().__init__(*args, **kwargs)
-        self.source = source
-
     def get_packages(self) -> set[str]:
         """Download and parse online source of top Python Package Index packages."""
         packages_info = self._download()
@@ -47,9 +40,7 @@ def _download(self) -> dict[str, Any]:
         except requests.exceptions.JSONDecodeError as err:
             raise InvalidJSONError from err
 
-        logger.debug(
-            f"Successfully downloaded trusted packages list from {self.source}"
-        )
+        logger.debug(f"Successfully downloaded trusted packages list from {self.source}")
 
         return packages_json
 

tests/core/test_config_handler.py

@@ -4,7 +4,7 @@
 
 import pytest
 from tomlkit import TOMLDocument, dumps, parse
-from twyn.base.constants import AvailableLoggingLevels
+from twyn.base.constants import DEFAULT_TOP_PYPI_PACKAGES, AvailableLoggingLevels
 from twyn.core.config_handler import ConfigHandler, ReadTwynConfiguration, TwynConfiguration
 from twyn.core.exceptions import (
     AllowlistPackageAlreadyExistsError,
@@ -30,7 +30,11 @@ def test_no_enforce_file_on_non_existent_file(self, mock_is_file):
         config = ConfigHandler(enforce_file=False).resolve_config()
 
         assert config == TwynConfiguration(
-            dependency_file=None, selector_method="all", logging_level=AvailableLoggingLevels.warning, allowlist=set()
+            dependency_file=None,
+            selector_method="all",
+            logging_level=AvailableLoggingLevels.warning,
+            allowlist=set(),
+            pypi_reference=DEFAULT_TOP_PYPI_PACKAGES,
         )
 
     def test_config_raises_for_unknown_file(self):
@@ -54,6 +58,7 @@ def test_get_twyn_data_from_file(self, pyproject_toml_file):
             selector_method="my_selector",
             logging_level="debug",
             allowlist={"boto4", "boto2"},
+            pypi_reference=None,
         )
 
     def test_write_toml(self, pyproject_toml_file):
@@ -91,6 +96,7 @@ def test_write_toml(self, pyproject_toml_file):
                     "selector_method": "my_selector",
                     "logging_level": "debug",
                     "allowlist": {},
+                    "pypi_reference": DEFAULT_TOP_PYPI_PACKAGES,
                 },
             }
         }

tests/main/test_main.py

@@ -1,7 +1,7 @@
 from unittest.mock import patch
 
 import pytest
-from twyn.base.constants import AvailableLoggingLevels
+from twyn.base.constants import DEFAULT_TOP_PYPI_PACKAGES, AvailableLoggingLevels
 from twyn.core.config_handler import ConfigHandler, TwynConfiguration, _get_logging_level
 from twyn.dependency_parser import RequirementsTxtParser
 from twyn.main import (
@@ -34,6 +34,7 @@ class TestCheckDependencies:
                     selector_method="first-letter",
                     logging_level=AvailableLoggingLevels.info,
                     allowlist={"boto4", "boto2"},
+                    pypi_reference=DEFAULT_TOP_PYPI_PACKAGES,
                 ),
             ),
             (
@@ -51,6 +52,7 @@ class TestCheckDependencies:
                     selector_method="nearby-letter",
                     logging_level=AvailableLoggingLevels.debug,
                     allowlist={"boto4", "boto2"},
+                    pypi_reference=DEFAULT_TOP_PYPI_PACKAGES,
                 ),
             ),
             (
@@ -63,6 +65,7 @@ class TestCheckDependencies:
                     selector_method="all",
                     logging_level=AvailableLoggingLevels.warning,
                     allowlist=set(),
+                    pypi_reference=DEFAULT_TOP_PYPI_PACKAGES,
                 ),
             ),
             (
@@ -81,6 +84,7 @@ class TestCheckDependencies:
                     selector_method=None,
                     logging_level=AvailableLoggingLevels.debug,
                     allowlist=set(),
+                    pypi_reference=DEFAULT_TOP_PYPI_PACKAGES,
                 ),
             ),
         ],
@@ -232,6 +236,7 @@ def test_check_dependencies_ignores_package_in_allowlist(
             dependency_file=None,
             selector_method="first-letter",
             logging_level=AvailableLoggingLevels.info,
+            pypi_reference=DEFAULT_TOP_PYPI_PACKAGES,
         )
 
         with patch("twyn.main.ConfigHandler.resolve_config", return_value=m_config):

tests/trusted_packages/test_references.py

@@ -35,7 +35,7 @@ def get_packages(self) -> set[str]:
             return {"foo", "bar"}
 
     def test_get_packages(self):
-        assert self.HardcodedPackageReference().get_packages() == {"foo", "bar"}
+        assert self.HardcodedPackageReference(source="foo").get_packages() == {"foo", "bar"}
 
 
 class TestTopPyPiReference: