diff --git a/src/strands_tools/bright_data.py b/src/strands_tools/bright_data.py index 8460c634..ac0cfd9f 100644 --- a/src/strands_tools/bright_data.py +++ b/src/strands_tools/bright_data.py @@ -73,6 +73,7 @@ import logging import os import time +from pathlib import Path from typing import Dict, Optional from urllib.parse import quote @@ -88,6 +89,48 @@ console = console_util.create() +def _resolve_output_path(output_path: str) -> Path: + """Resolve an output path and confine it to an allowed root directory. + + The allowed root is the directory named by the STRANDS_BRIGHT_DATA_OUTPUT_DIR + environment variable, or the current working directory when that variable is + not set. The resolved path must be the root itself or located inside it, and + the final path component must not be a symlink that points elsewhere. + + Args: + output_path (str): User-supplied path to save the screenshot. + + Returns: + Path: The resolved, confined path. + + Raises: + ValueError: If the resolved path falls outside the allowed root. + """ + root = Path(os.environ.get("STRANDS_BRIGHT_DATA_OUTPUT_DIR", Path.cwd())).expanduser().resolve() + + candidate = Path(output_path).expanduser() + if not candidate.is_absolute(): + candidate = root / candidate + + # Reject a final component that is a symlink so we do not follow it elsewhere. + if candidate.is_symlink(): + raise ValueError( + f"output_path must not be a symlink: {output_path}. output_path must be within the " + f"working directory ({root}) or the directory named by STRANDS_BRIGHT_DATA_OUTPUT_DIR." + ) + + resolved = candidate.resolve() + + if resolved != root and root not in resolved.parents: + raise ValueError( + f"output_path must be within the working directory ({root}) or the directory named by " + f"STRANDS_BRIGHT_DATA_OUTPUT_DIR. Set STRANDS_BRIGHT_DATA_OUTPUT_DIR to choose a different " + f"directory." + ) + + return resolved + + class BrightDataClient: """Client for interacting with Bright Data API.""" @@ -158,13 +201,25 @@ def get_screenshot(self, url: str, output_path: str, zone: Optional[str] = None) """ Take a screenshot of a webpage. + The screenshot is written within an allowed output directory. By default + this is the current working directory; set the STRANDS_BRIGHT_DATA_OUTPUT_DIR + environment variable to write somewhere else. Paths that resolve outside the + allowed directory (for example via ".." traversal, an absolute path outside + it, or a symlinked final component) are rejected. + Args: url (str): URL to screenshot - output_path (str): Path to save the screenshot + output_path (str): Path to save the screenshot. Relative paths are taken + from the allowed output directory. The resolved path must stay within + that directory (the current working directory by default, or + STRANDS_BRIGHT_DATA_OUTPUT_DIR when set). zone (Optional[str]): Override default zone Returns: str: Path to saved screenshot + + Raises: + ValueError: If output_path resolves outside the allowed output directory. """ payload = {"url": url, "zone": zone or self.zone, "format": "raw", "data_format": "screenshot"} @@ -173,10 +228,12 @@ def get_screenshot(self, url: str, output_path: str, zone: Optional[str] = None) if response.status_code != 200: raise Exception(f"Error {response.status_code}: {response.text}") - with open(output_path, "wb") as f: + resolved_path = _resolve_output_path(output_path) + + with open(resolved_path, "wb") as f: f.write(response.content) - return output_path + return str(resolved_path) @staticmethod def encode_query(query: str) -> str: @@ -421,7 +478,11 @@ def bright_data( Args: action: The action to perform (scrape_as_markdown, get_screenshot, search_engine, web_data_feed) url: URL to scrape or extract data from (for scrape_as_markdown, get_screenshot, web_data_feed) - output_path: Path to save the screenshot (for get_screenshot) + output_path: Path to save the screenshot (for get_screenshot). The screenshot is written + within an allowed output directory: the current working directory by default, or the + directory named by the STRANDS_BRIGHT_DATA_OUTPUT_DIR environment variable when set. + STRANDS_BRIGHT_DATA_OUTPUT_DIR may point anywhere the operator chooses. Relative paths + are taken from that directory; paths that resolve outside it are rejected. zone: Web Unlocker zone name (optional). If not provided, uses BRIGHTDATA_ZONE environment variable, or defaults to "web_unlocker1". Set BRIGHTDATA_ZONE in your .env file to configure your specific Web Unlocker zone name (e.g., BRIGHTDATA_ZONE=web_unlocker_12345) diff --git a/tests/test_bright_data.py b/tests/test_bright_data.py index d5633991..7fb6e16b 100644 --- a/tests/test_bright_data.py +++ b/tests/test_bright_data.py @@ -311,3 +311,115 @@ def test_zone_default_fallback(): # Verify the client was created with the default zone mock_client_class.assert_called_with(verbose=True, zone="web_unlocker1") + + +def _mock_screenshot_response(): + """Return a mocked successful screenshot HTTP response.""" + response = MagicMock() + response.status_code = 200 + response.content = b"\x89PNG fake screenshot bytes" + return response + + +@patch.dict(os.environ, {"BRIGHTDATA_API_KEY": "test_api_key"}) +def test_get_screenshot_writes_within_output_dir(tmp_path): + """A relative output_path is written inside the configured output directory.""" + with patch.dict(os.environ, {"STRANDS_BRIGHT_DATA_OUTPUT_DIR": str(tmp_path)}): + with patch("strands_tools.bright_data.requests.post", return_value=_mock_screenshot_response()): + client = BrightDataClient() + result = client.get_screenshot("https://example.com", "shot.png") + + expected = tmp_path / "shot.png" + assert result == str(expected) + assert expected.read_bytes() == b"\x89PNG fake screenshot bytes" + + +@patch.dict(os.environ, {"BRIGHTDATA_API_KEY": "test_api_key"}) +def test_get_screenshot_rejects_relative_traversal(tmp_path): + """A relative path escaping the output directory is rejected and nothing is written.""" + with patch.dict(os.environ, {"STRANDS_BRIGHT_DATA_OUTPUT_DIR": str(tmp_path)}): + with patch("strands_tools.bright_data.requests.post", return_value=_mock_screenshot_response()): + client = BrightDataClient() + with pytest.raises(ValueError): + client.get_screenshot("https://example.com", "../../etc/x") + + +@patch.dict(os.environ, {"BRIGHTDATA_API_KEY": "test_api_key"}) +def test_get_screenshot_rejects_absolute_path_outside_root(tmp_path): + """An absolute path outside the output directory is rejected.""" + with patch.dict(os.environ, {"STRANDS_BRIGHT_DATA_OUTPUT_DIR": str(tmp_path)}): + target = "/tmp/strands_bright_data_should_not_exist.png" + with patch("strands_tools.bright_data.requests.post", return_value=_mock_screenshot_response()): + client = BrightDataClient() + with pytest.raises(ValueError): + client.get_screenshot("https://example.com", target) + + +@patch.dict(os.environ, {"BRIGHTDATA_API_KEY": "test_api_key"}) +def test_get_screenshot_rejects_symlink_target(tmp_path): + """A final path component that is a symlink is rejected.""" + outside = tmp_path / "outside_target.png" + root = tmp_path / "root" + root.mkdir() + link = root / "link.png" + link.symlink_to(outside) + + with patch.dict(os.environ, {"STRANDS_BRIGHT_DATA_OUTPUT_DIR": str(root)}): + with patch("strands_tools.bright_data.requests.post", return_value=_mock_screenshot_response()): + client = BrightDataClient() + with pytest.raises(ValueError): + client.get_screenshot("https://example.com", "link.png") + + assert not outside.exists() + + +@patch.dict(os.environ, {"BRIGHTDATA_API_KEY": "test_api_key"}) +def test_get_screenshot_default_root_rejects_traversal(tmp_path, monkeypatch): + """With no STRANDS_BRIGHT_DATA_OUTPUT_DIR set, confinement defaults to the working directory.""" + monkeypatch.delenv("STRANDS_BRIGHT_DATA_OUTPUT_DIR", raising=False) + monkeypatch.chdir(tmp_path) + + with patch("strands_tools.bright_data.requests.post", return_value=_mock_screenshot_response()): + client = BrightDataClient() + with pytest.raises(ValueError): + client.get_screenshot("https://example.com", "../../etc/x") + with pytest.raises(ValueError): + client.get_screenshot("https://example.com", "/tmp/strands_default_root_outside.png") + + assert not (tmp_path.parent.parent / "etc" / "x").exists() + + +@patch.dict(os.environ, {"BRIGHTDATA_API_KEY": "test_api_key"}) +def test_get_screenshot_rejection_message_is_actionable(tmp_path, monkeypatch): + """The rejection message points at the working directory and the env override.""" + monkeypatch.delenv("STRANDS_BRIGHT_DATA_OUTPUT_DIR", raising=False) + monkeypatch.chdir(tmp_path) + + with patch("strands_tools.bright_data.requests.post", return_value=_mock_screenshot_response()): + client = BrightDataClient() + with pytest.raises(ValueError) as excinfo: + client.get_screenshot("https://example.com", "/tmp/strands_actionable_message.png") + + message = str(excinfo.value) + assert "working directory" in message + assert "STRANDS_BRIGHT_DATA_OUTPUT_DIR" in message + + +@patch.dict(os.environ, {"BRIGHTDATA_API_KEY": "test_api_key"}) +def test_get_screenshot_env_dir_allows_arbitrary_operator_location(tmp_path, monkeypatch): + """An operator-chosen STRANDS_BRIGHT_DATA_OUTPUT_DIR allows writing within it, even via absolute path.""" + cwd = tmp_path / "cwd" + cwd.mkdir() + monkeypatch.chdir(cwd) + output_dir = tmp_path / "operator_chosen" + output_dir.mkdir() + + with patch.dict(os.environ, {"STRANDS_BRIGHT_DATA_OUTPUT_DIR": str(output_dir)}): + with patch("strands_tools.bright_data.requests.post", return_value=_mock_screenshot_response()): + client = BrightDataClient() + absolute_target = str(output_dir / "shot.png") + result = client.get_screenshot("https://example.com", absolute_target) + + expected = output_dir / "shot.png" + assert result == str(expected) + assert expected.read_bytes() == b"\x89PNG fake screenshot bytes"