diff --git a/obliteratus/ui_watchtower.py b/obliteratus/ui_watchtower.py index 1dffb05..75267f2 100644 --- a/obliteratus/ui_watchtower.py +++ b/obliteratus/ui_watchtower.py @@ -167,11 +167,24 @@ def _download_result(model_id: str): if not model_id: return None safe_name = model_id.strip().replace("/", "_").replace("\\", "_") - base = Path.home() / ".obliteratus" / "auto_obliterate" / safe_name + if not safe_name or safe_name in {".", ".."}: + return None + managed_root = Path.home() / ".obliteratus" / "auto_obliterate" + base = managed_root / safe_name # Find the latest iteration output if base.exists(): - base_resolved = base.resolve() + try: + managed_root_resolved = managed_root.resolve(strict=True) + base_resolved = base.resolve(strict=True) + except (OSError, ValueError): + return None + if ( + base_resolved == managed_root_resolved + or not base_resolved.is_relative_to(managed_root_resolved) + or not base_resolved.is_dir() + ): + return None def iteration_number(path: Path) -> int: try: @@ -179,7 +192,7 @@ def _download_result(model_id: str): except ValueError: return -1 - iter_dirs = sorted(base.glob("iter_*"), key=iteration_number, reverse=True) + iter_dirs = sorted(base_resolved.glob("iter_*"), key=iteration_number, reverse=True) for d in iter_dirs: try: resolved = d.resolve(strict=True) @@ -187,8 +200,8 @@ def _download_result(model_id: str): continue if not resolved.is_relative_to(base_resolved): continue - if d.is_dir() and (d / "config.json").is_file(): - return str(d) + if resolved.is_dir() and (resolved / "config.json").is_file(): + return str(resolved) return None diff --git a/tests/test_ui_watchtower_contracts.py b/tests/test_ui_watchtower_contracts.py index 40523f4..fe81bd6 100644 --- a/tests/test_ui_watchtower_contracts.py +++ b/tests/test_ui_watchtower_contracts.py @@ -138,6 +138,22 @@ def test_download_result_handles_empty_id_and_unreadable_iteration(monkeypatch, assert ui_watchtower._download_result("org/model") == str(valid) +def test_download_result_handles_unreadable_managed_model_directory(monkeypatch, tmp_path): + monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) + base = tmp_path / ".obliteratus" / "auto_obliterate" / "org_model" + base.mkdir(parents=True) + resolve = Path.resolve + + def controlled_resolve(path, *args, **kwargs): + if path == base: + raise OSError("unreadable model directory") + return resolve(path, *args, **kwargs) + + monkeypatch.setattr(Path, "resolve", controlled_resolve) + + assert ui_watchtower._download_result("org/model") is None + + def test_download_result_rejects_symlink_that_escapes_managed_root(monkeypatch, tmp_path): monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) base = tmp_path / ".obliteratus" / "auto_obliterate" / "org_model" @@ -153,6 +169,30 @@ def test_download_result_rejects_symlink_that_escapes_managed_root(monkeypatch, assert ui_watchtower._download_result("org/model") == str(valid) +@pytest.mark.parametrize("model_id", [".", ".."]) +def test_download_result_rejects_managed_root_aliases(monkeypatch, tmp_path, model_id): + monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) + root = tmp_path / ".obliteratus" / "auto_obliterate" + iteration = root / "iter_10" + iteration.mkdir(parents=True) + (iteration / "config.json").write_text("{}", encoding="utf-8") + + assert ui_watchtower._download_result(model_id) is None + + +def test_download_result_rejects_model_directory_symlink_escape(monkeypatch, tmp_path): + monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path)) + root = tmp_path / ".obliteratus" / "auto_obliterate" + root.mkdir(parents=True) + outside = tmp_path / "outside" + iteration = outside / "iter_99" + iteration.mkdir(parents=True) + (iteration / "config.json").write_text("{}", encoding="utf-8") + (root / "org_model").symlink_to(outside, target_is_directory=True) + + assert ui_watchtower._download_result("org/model") is None + + def test_scan_handler_returns_data_and_escaped_failure(monkeypatch): watchtower = SimpleNamespace( scan=lambda on_log: (on_log("scanned") or [SimpleNamespace()]),