fix(bright_data): confine get_screenshot output to a configurable directory

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)

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"