diff --git a/src/twyn/cli.py b/src/twyn/cli.py index 1e69517..7af1948 100644 --- a/src/twyn/cli.py +++ b/src/twyn/cli.py @@ -111,6 +111,16 @@ def entry_point() -> None: default=False, help="Recursively look for files when trying to locate them automatically. Ignored if --dependency-file is given.", ) +@click.option( + "--pypi-source", + type=str, + help="Alternative PyPI source URL to use for fetching trusted packages.", +) +@click.option( + "--npm-source", + type=str, + help="Alternative npm source URL to use for fetching trusted packages.", +) def run( # noqa: C901 config: str, dependency_file: tuple[str], @@ -123,6 +133,8 @@ def run( # noqa: C901 json: bool, package_ecosystem: Optional[str], recursive: bool, + pypi_source: Optional[str], + npm_source: Optional[str], ) -> int: if vv: logger.setLevel(logging.DEBUG) @@ -149,6 +161,8 @@ def run( # noqa: C901 load_config_from_file=True, package_ecosystem=package_ecosystem, recursive=recursive, + pypi_source=pypi_source, + npm_source=npm_source, ) except TwynError as e: raise CliError(str(e)) from e diff --git a/src/twyn/config/config_handler.py b/src/twyn/config/config_handler.py index 522ed11..9c03966 100644 --- a/src/twyn/config/config_handler.py +++ b/src/twyn/config/config_handler.py @@ -62,13 +62,15 @@ class ConfigHandler: def __init__(self, file_handler: Optional[FileHandler] = None) -> None: self.file_handler = file_handler - def resolve_config( + def resolve_config( # noqa: C901, PLR0912 self, selector_method: Optional[str] = None, dependency_files: Optional[set[str]] = None, use_cache: Optional[bool] = None, package_ecosystem: Optional[PackageEcosystems] = None, recursive: Optional[bool] = None, + pypi_source: Optional[str] = None, + npm_source: Optional[str] = None, ) -> TwynConfiguration: """Resolve the configuration for Twyn. @@ -107,12 +109,28 @@ def resolve_config( else: final_recursive = DEFAULT_RECURSIVE + # Determine final pypi_source from CLI, config file, or default + if pypi_source is not None: + final_pypi_source = pypi_source + elif read_config.pypi_source is not None: + final_pypi_source = read_config.pypi_source + else: + final_pypi_source = None + + # Determine final npm_source from CLI, config file, or default + if npm_source is not None: + final_npm_source = npm_source + elif read_config.npm_source is not None: + final_npm_source = read_config.npm_source + else: + final_npm_source = None + return TwynConfiguration( dependency_files=dependency_files or read_config.dependency_files or set(), selector_method=final_selector_method, allowlist=read_config.allowlist, - pypi_source=read_config.pypi_source, - npm_source=read_config.npm_source, + pypi_source=final_pypi_source, + npm_source=final_npm_source, use_cache=final_use_cache, package_ecosystem=package_ecosystem or read_config.package_ecosystem, recursive=final_recursive, diff --git a/src/twyn/main.py b/src/twyn/main.py index 8d4a042..e9f118e 100644 --- a/src/twyn/main.py +++ b/src/twyn/main.py @@ -44,6 +44,8 @@ def check_dependencies( load_config_from_file: bool = False, package_ecosystem: Optional[PackageEcosystems] = None, recursive: Optional[bool] = None, + pypi_source: Optional[str] = None, + npm_source: Optional[str] = None, ) -> TyposquatCheckResults: """ Check if the provided dependencies are potential typosquats of trusted packages. @@ -72,6 +74,8 @@ def check_dependencies( use_cache=use_cache, package_ecosystem=package_ecosystem, recursive=recursive, + pypi_source=pypi_source, + npm_source=npm_source, ) maybe_cache_handler = CacheHandler() if config.use_cache else None selector_method_obj = _get_selector_method(config.selector_method) @@ -284,6 +288,8 @@ def _get_config( use_cache: Optional[bool], package_ecosystem: Optional[PackageEcosystems], recursive: Optional[bool], + pypi_source: Optional[str], + npm_source: Optional[str], ) -> TwynConfiguration: """Given the arguments passed to the main function and the configuration loaded from the config file (if any), return a config object.""" if load_config_from_file: @@ -296,4 +302,6 @@ def _get_config( use_cache=use_cache, package_ecosystem=package_ecosystem, recursive=recursive, + pypi_source=pypi_source, + npm_source=npm_source, ) diff --git a/tests/main/test_cli.py b/tests/main/test_cli.py index fae0f57..ab72341 100644 --- a/tests/main/test_cli.py +++ b/tests/main/test_cli.py @@ -96,6 +96,8 @@ def test_click_arguments_dependency_file(self, mock_check_dependencies: Mock) -> load_config_from_file=True, package_ecosystem=None, recursive=False, + pypi_source=None, + npm_source=None, ) ] @@ -121,6 +123,8 @@ def test_click_arguments_dependency_file_in_different_path(self, mock_check_depe load_config_from_file=True, package_ecosystem=None, recursive=False, + pypi_source=None, + npm_source=None, ) ] @@ -146,6 +150,8 @@ def test_package_ecosystem_option(self, mock_check_dependencies: Mock) -> None: load_config_from_file=True, package_ecosystem="pypi", recursive=False, + pypi_source=None, + npm_source=None, ) ] @@ -170,6 +176,8 @@ def test_click_arguments_single_dependency_cli(self, mock_check_dependencies: Mo load_config_from_file=True, package_ecosystem=None, recursive=False, + pypi_source=None, + npm_source=None, ) ] @@ -206,6 +214,8 @@ def test_click_arguments_multiple_dependencies(self, mock_check_dependencies: Mo load_config_from_file=True, package_ecosystem=None, recursive=False, + pypi_source=None, + npm_source=None, ) ] @@ -235,6 +245,8 @@ def test_recursive(self, mock_check_dependencies: Mock) -> None: load_config_from_file=True, package_ecosystem=None, recursive=True, + pypi_source=None, + npm_source=None, ) assert mock_check_dependencies.call_args_list[0] == call_args assert mock_check_dependencies.call_args_list[1] == call_args @@ -255,6 +267,8 @@ def test_click_arguments_default(self, mock_check_dependencies: Mock) -> None: load_config_from_file=True, package_ecosystem=None, recursive=False, + pypi_source=None, + npm_source=None, ) ] @@ -387,3 +401,89 @@ def test_unhandled_exception_is_caught_and_wrapped_in_cli_error( assert isinstance(result.exception, SystemExit) # Check that the generic error message was logged assert "Unhandled exception occured." in caplog.text + + @patch("twyn.cli.check_dependencies") + def test_pypi_source_option(self, mock_check_dependencies: Mock) -> None: + """Test that --pypi-source option is passed correctly.""" + runner = CliRunner() + runner.invoke( + cli.run, + [ + "--pypi-source", + "https://custom-pypi.org/", + ], + ) + + assert mock_check_dependencies.call_args_list == [ + call( + config_file=None, + dependency_files=None, + dependencies=None, + selector_method=None, + use_cache=None, + show_progress_bar=True, + load_config_from_file=True, + package_ecosystem=None, + recursive=False, + pypi_source="https://custom-pypi.org/", + npm_source=None, + ) + ] + + @patch("twyn.cli.check_dependencies") + def test_npm_source_option(self, mock_check_dependencies: Mock) -> None: + """Test that --npm-source option is passed correctly.""" + runner = CliRunner() + runner.invoke( + cli.run, + [ + "--npm-source", + "https://custom-npm.org/", + ], + ) + + assert mock_check_dependencies.call_args_list == [ + call( + config_file=None, + dependency_files=None, + dependencies=None, + selector_method=None, + use_cache=None, + show_progress_bar=True, + load_config_from_file=True, + package_ecosystem=None, + recursive=False, + pypi_source=None, + npm_source="https://custom-npm.org/", + ) + ] + + @patch("twyn.cli.check_dependencies") + def test_both_source_options(self, mock_check_dependencies: Mock) -> None: + """Test that both --pypi-source and --npm-source options work together.""" + runner = CliRunner() + runner.invoke( + cli.run, + [ + "--pypi-source", + "https://custom-pypi.org/", + "--npm-source", + "https://custom-npm.org/", + ], + ) + + assert mock_check_dependencies.call_args_list == [ + call( + config_file=None, + dependency_files=None, + dependencies=None, + selector_method=None, + use_cache=None, + show_progress_bar=True, + load_config_from_file=True, + package_ecosystem=None, + recursive=False, + pypi_source="https://custom-pypi.org/", + npm_source="https://custom-npm.org/", + ) + ] diff --git a/tests/main/test_main.py b/tests/main/test_main.py index 2101edc..f8a59bc 100644 --- a/tests/main/test_main.py +++ b/tests/main/test_main.py @@ -50,6 +50,7 @@ class TestCheckDependencies: "recursive": False, "pypi_source": "a", "npm_source": "a", + "package_ecosystem": "npm", }, TwynConfiguration( dependency_files={"requirements.txt"}, @@ -72,6 +73,7 @@ class TestCheckDependencies: "recursive": True, "pypi_source": "pypi", "npm_source": "npm", + "package_ecosystem": "pypi", }, TwynConfiguration( dependency_files={"poetry.lock"}, @@ -80,7 +82,7 @@ class TestCheckDependencies: pypi_source="pypi", npm_source="npm", use_cache=False, - package_ecosystem=None, + package_ecosystem="pypi", recursive=True, ), ), # Config from file takes precendence over fallback values @@ -122,6 +124,8 @@ def test_options_priorities_assignation( selector_method=cli_config.get("selector_method"), dependency_files=cli_config.get("dependency_file", set()), use_cache=cli_config.get("use_cache"), + pypi_source=cli_config.get("pypi_source"), + npm_source=cli_config.get("npm_source"), recursive=cli_config.get("recursive"), package_ecosystem=cli_config.get("package_ecosystem"), ) @@ -129,9 +133,10 @@ def test_options_priorities_assignation( assert resolved.dependency_files == expected_resolved_config.dependency_files assert resolved.selector_method == expected_resolved_config.selector_method assert resolved.allowlist == expected_resolved_config.allowlist - assert resolved.use_cache == expected_resolved_config.use_cache assert resolved.package_ecosystem == expected_resolved_config.package_ecosystem assert resolved.recursive == expected_resolved_config.recursive + assert resolved.npm_source == expected_resolved_config.npm_source + assert resolved.pypi_source == expected_resolved_config.pypi_source @patch("twyn.trusted_packages.TopPyPiReference.get_packages") def test_check_dependencies_detects_typosquats_from_file(