From 68f479f8654f53b2af5085caa64cefb7b8924c17 Mon Sep 17 00:00:00 2001 From: DevilsNerve Date: Mon, 20 Jul 2026 16:48:37 -0500 Subject: [PATCH] Harden directory input selection --- pkzip_overdrive.py | 68 +++++++++++++++++++++++++---------- tests/test_pkzip_overdrive.py | 36 +++++++++++++++++-- 2 files changed, 83 insertions(+), 21 deletions(-) diff --git a/pkzip_overdrive.py b/pkzip_overdrive.py index e7605d4..de8bc4d 100644 --- a/pkzip_overdrive.py +++ b/pkzip_overdrive.py @@ -667,13 +667,24 @@ def choose_source_path(provided: str | None) -> Path: provided = provided.strip() if len(provided) >= 2 and provided[0] == provided[-1] and provided[0] in "\"'": provided = provided[1:-1] + if not provided: + raise SystemExit("error: ZIP/hash path cannot be empty") source = Path(os.path.expandvars(provided)).expanduser().resolve() if source.is_dir(): - archives = sorted( - (path for path in source.glob("*.zip") if path.is_file()), - key=lambda path: path.name.casefold(), - ) + try: + archives = sorted( + ( + path + for path in source.iterdir() + if path.is_file() and path.suffix.casefold() == ".zip" + ), + key=lambda path: path.name.casefold(), + ) + except OSError as error: + raise SystemExit( + f"error: cannot list ZIP directory {source}: {error}" + ) from error if not archives: raise SystemExit(f"error: no ZIP files found directly in directory: {source}") if len(archives) == 1: @@ -686,12 +697,17 @@ def choose_source_path(provided: str | None) -> Path: while True: try: selected = input(f"Select ZIP (1-{len(archives)}): ").strip() + except EOFError as error: + raise SystemExit( + "error: interactive input ended; provide an explicit ZIP file path" + ) from error + try: selected_index = int(selected) - if 1 <= selected_index <= len(archives): - source = archives[selected_index - 1] - break - except (EOFError, ValueError): - pass + except ValueError: + selected_index = 0 + if 1 <= selected_index <= len(archives): + source = archives[selected_index - 1] + break print("error: enter one of the displayed numbers", file=sys.stderr) if not source.is_file(): @@ -740,15 +756,24 @@ def choose_wordlist_path(provided: str | None) -> Path: provided = provided.strip() if len(provided) >= 2 and provided[0] == provided[-1] and provided[0] in "\"'": provided = provided[1:-1] + if not provided: + raise SystemExit("error: wordlist path cannot be empty") wordlist = Path(os.path.expandvars(provided)).expanduser().resolve() if wordlist.is_dir(): preferred_suffixes = {".txt", ".dict", ".dic", ".lst", ".wordlist", ".gz"} - all_files = sorted( - (path for path in wordlist.iterdir() if path.is_file()), - key=lambda path: path.name.casefold(), - ) - preferred = [path for path in all_files if path.suffix.casefold() in preferred_suffixes] + try: + all_files = sorted( + (path for path in wordlist.iterdir() if path.is_file()), + key=lambda path: path.name.casefold(), + ) + except OSError as error: + raise SystemExit( + f"error: cannot list wordlist directory {wordlist}: {error}" + ) from error + preferred = [ + path for path in all_files if path.suffix.casefold() in preferred_suffixes + ] files = preferred or all_files if not files: raise SystemExit(f"error: no wordlist files found in directory: {wordlist}") @@ -762,12 +787,17 @@ def choose_wordlist_path(provided: str | None) -> Path: while True: try: selected = input(f"Select wordlist (1-{len(files)}): ").strip() + except EOFError as error: + raise SystemExit( + "error: interactive input ended; provide an explicit wordlist file path" + ) from error + try: selected_index = int(selected) - if 1 <= selected_index <= len(files): - wordlist = files[selected_index - 1] - break - except (EOFError, ValueError): - pass + except ValueError: + selected_index = 0 + if 1 <= selected_index <= len(files): + wordlist = files[selected_index - 1] + break print("error: enter one of the displayed numbers", file=sys.stderr) if not wordlist.is_file(): diff --git a/tests/test_pkzip_overdrive.py b/tests/test_pkzip_overdrive.py index 4de2e0b..9bece1b 100644 --- a/tests/test_pkzip_overdrive.py +++ b/tests/test_pkzip_overdrive.py @@ -148,15 +148,26 @@ def test_extract_pkzip_hash_writes_only_the_record(self) -> None: mode = stat.S_IMODE(destination.stat().st_mode) self.assertEqual(mode, 0o600) - def test_choose_source_path_uses_a_directorys_only_zip(self) -> None: + def test_choose_source_path_uses_a_directorys_only_zip_case_insensitively(self) -> None: with tempfile.TemporaryDirectory() as directory_name: directory = Path(directory_name) - expected = directory / "archive.zip" + expected = directory / "archive.ZIP" expected.write_bytes(b"not needed for path selection") with mock.patch("builtins.print"): selected = overdrive.choose_source_path(str(directory)) self.assertEqual(selected, expected.resolve()) + def test_multiple_archives_fail_cleanly_when_interactive_input_ends(self) -> None: + with tempfile.TemporaryDirectory() as directory_name: + directory = Path(directory_name) + (directory / "first.zip").write_bytes(b"first") + (directory / "second.zip").write_bytes(b"second") + with mock.patch("builtins.print"), mock.patch( + "builtins.input", side_effect=EOFError + ): + with self.assertRaisesRegex(SystemExit, "explicit ZIP file path"): + overdrive.choose_source_path(str(directory)) + def test_choose_wordlist_prefers_known_wordlist_suffixes(self) -> None: with tempfile.TemporaryDirectory() as directory_name: directory = Path(directory_name) @@ -167,6 +178,27 @@ def test_choose_wordlist_prefers_known_wordlist_suffixes(self) -> None: selected = overdrive.choose_wordlist_path(str(directory)) self.assertEqual(selected, expected.resolve()) + def test_multiple_wordlists_fail_cleanly_when_interactive_input_ends(self) -> None: + with tempfile.TemporaryDirectory() as directory_name: + directory = Path(directory_name) + (directory / "first.txt").write_text("first\n", encoding="utf-8") + (directory / "second.txt").write_text("second\n", encoding="utf-8") + with mock.patch("builtins.print"), mock.patch( + "builtins.input", side_effect=EOFError + ): + with self.assertRaisesRegex(SystemExit, "explicit wordlist file path"): + overdrive.choose_wordlist_path(str(directory)) + + def test_empty_paths_are_rejected(self) -> None: + for chooser in ( + overdrive.choose_source_path, + overdrive.choose_wordlist_path, + ): + for value in ("", '""', "''"): + with self.subTest(chooser=chooser.__name__, value=value): + with self.assertRaisesRegex(SystemExit, "path cannot be empty"): + chooser(value) + def test_identify_mode_requires_one_supported_match(self) -> None: completed = subprocess.CompletedProcess( args=[],