Feature/custom search

cloakbrowser/browser.py

@@ -14,8 +14,10 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import os
+from pathlib import Path
 from typing import Any, Literal, TypedDict
 from urllib.parse import quote, unquote, urlparse, urlunparse
 
@@ -29,6 +31,66 @@
 # Sentinel to distinguish "viewport not provided" from "viewport=None" (disable emulation)
 _VIEWPORT_UNSET = object()
 
+_SEARCH_ENGINE_PREFS: dict[str, dict] = {
+    "google": {
+        "short_name": "Google",
+        "keyword": "google.com",
+        "url": "https://www.google.com/search?q={searchTerms}",
+        "new_tab_url": "https://www.google.com/_/chrome/newtab",
+        "suggestions_url": "https://www.google.com/complete/search?client=chrome&q={searchTerms}",
+        "favicon_url": "https://www.google.com/favicon.ico",
+        "input_encodings": ["UTF-8"],
+        "safe_for_autoreplace": True,
+    },
+    "bing": {
+        "short_name": "Bing",
+        "keyword": "bing.com",
+        "url": "http://www.bing.com/search?setmkt=en-US&q={searchTerms}",
+        "new_tab_url": "https://www.bing.com/chrome/newtab",
+        "suggestions_url": "http://api.bing.com/osjson.aspx?query={searchTerms}&language={language}",
+        "favicon_url": "http://www.bing.com/s/wlflag.ico",
+        "input_encodings": ["UTF-8"],
+        "prepopulate_id": 3,
+        "safe_for_autoreplace": True,
+    },
+    "duckduckgo": {
+        "short_name": "DuckDuckGo",
+        "keyword": "duckduckgo.com",
+        "url": "https://duckduckgo.com/?q={searchTerms}&ia=web",
+        "new_tab_url": "",
+        "suggestions_url": "https://duckduckgo.com/ac/?q={searchTerms}&type=list",
+        "favicon_url": "https://duckduckgo.com/favicon.ico",
+        "input_encodings": ["UTF-8"],
+        "safe_for_autoreplace": True,
+    },
+}
+
+
+def _apply_search_engine_prefs(
+    user_data_dir: str | os.PathLike,
+    search_engine: str,
+) -> None:
+    """Write default_search_provider_data into Chromium's Default/Preferences file."""
+    prefs_data = _SEARCH_ENGINE_PREFS.get(search_engine)
+    if not prefs_data:
+        logger.warning("Unknown search_engine %r - must be 'google', 'bing', or 'duckduckgo'.", search_engine)
+        return
+
+    profile_dir = Path(os.fspath(user_data_dir)) / "Default"
+    profile_dir.mkdir(parents=True, exist_ok=True)
+    prefs_file = profile_dir / "Preferences"
+
+    existing: dict = {}
+    if prefs_file.exists():
+        try:
+            existing = json.loads(prefs_file.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, OSError):
+            existing = {}
+
+    existing.setdefault("default_search_provider_data", {})["template_url_data"] = prefs_data
+    prefs_file.write_text(json.dumps(existing), encoding="utf-8")
+    logger.debug("Applied search_engine=%r to %s", search_engine, prefs_file)
+
 
 def _resolve_timezone(timezone: str | None, kwargs: dict[str, Any]) -> str | None:
     """Accept both timezone and timezone_id — either works, no warning."""
@@ -253,6 +315,7 @@ def launch_persistent_context(
     locale: str | None = None,
     timezone: str | None = None,
     color_scheme: Literal["light", "dark", "no-preference"] | None = None,
+    search_engine: Literal["google", "bing", "duckduckgo"] | None = None,
     geoip: bool = False,
     backend: str | None = None,
     humanize: bool = False,
@@ -283,6 +346,7 @@ def launch_persistent_context(
         timezone: IANA timezone (e.g. 'America/New_York').
         color_scheme: Color scheme preference — 'light', 'dark', or 'no-preference'.
             Default: None (uses Chromium default, which is 'light').
+        search_engine: Default search engine — 'google', 'bing', or 'duckduckgo'.
         geoip: Auto-detect timezone/locale from proxy IP (default False).
             Requires ``pip install cloakbrowser[geoip]``.
         backend: Playwright backend — 'playwright' (default) or 'patchright'.
@@ -334,6 +398,8 @@ def launch_persistent_context(
         context_kwargs["viewport"] = viewport
     if color_scheme:
         context_kwargs["color_scheme"] = color_scheme
+    if search_engine:
+        _apply_search_engine_prefs(user_data_dir, search_engine)
     context_kwargs.update(kwargs)
 
     pw = sync_playwright().start()
@@ -379,6 +445,7 @@ async def launch_persistent_context_async(
     locale: str | None = None,
     timezone: str | None = None,
     color_scheme: Literal["light", "dark", "no-preference"] | None = None,
+    search_engine: Literal["google", "bing", "duckduckgo"] | None = None,
     geoip: bool = False,
     backend: str | None = None,
     humanize: bool = False,
@@ -407,6 +474,7 @@ async def launch_persistent_context_async(
         locale: Browser locale, e.g. "en-US".
         timezone: IANA timezone (e.g. 'America/New_York').
         color_scheme: Color scheme preference — 'light', 'dark', or 'no-preference'.
+        search_engine: Default search engine — 'google', 'bing', or 'duckduckgo'.
         geoip: Auto-detect timezone/locale from proxy IP (default False).
         backend: Playwright backend — 'playwright' (default) or 'patchright'.
         humanize: Enable human-like mouse, keyboard, scroll behavior (default False).
@@ -462,6 +530,8 @@ async def launch_persistent_context_async(
         context_kwargs["viewport"] = viewport
     if color_scheme:
         context_kwargs["color_scheme"] = color_scheme
+    if search_engine:
+        _apply_search_engine_prefs(user_data_dir, search_engine)
     context_kwargs.update(kwargs)
 
     pw = await async_playwright().start()

tests/test_persistent_context.py

@@ -3,6 +3,7 @@
 All tests mock playwright to avoid needing a binary.
 """
 
+import json
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -116,6 +117,28 @@ def test_persistent_context_color_scheme(_mock_geoip, _mock_bin):
     assert call_kwargs["color_scheme"] == "dark"
 
 
+@patch("cloakbrowser.browser.ensure_binary", return_value="/fake/chrome")
+@patch("cloakbrowser.browser.maybe_resolve_geoip", return_value=(None, None, None))
+def test_persistent_context_search_engine_writes_chromium_prefs(_mock_geoip, _mock_bin, tmp_path):
+    """search_engine writes Chromium Default/Preferences before launch."""
+    pw_cm, pw, context = _make_mock_pw_and_context()
+    profile_dir = tmp_path / "profile"
+
+    with patch("playwright.sync_api.sync_playwright", return_value=pw_cm):
+        from cloakbrowser.browser import launch_persistent_context
+        launch_persistent_context(profile_dir, search_engine="bing")
+
+    prefs_path = profile_dir / "Default" / "Preferences"
+    prefs = json.loads(prefs_path.read_text(encoding="utf-8"))
+    data = prefs["default_search_provider_data"]["template_url_data"]
+    assert data["short_name"] == "Bing"
+    assert data["keyword"] == "bing.com"
+    assert data["prepopulate_id"] == 3
+    assert data["url"] == "http://www.bing.com/search?setmkt=en-US&q={searchTerms}"
+    assert data["suggestions_url"] == "http://api.bing.com/osjson.aspx?query={searchTerms}&language={language}"
+    assert "suggest_url" not in data
+
+
 @patch("cloakbrowser.browser.maybe_resolve_geoip", return_value=("Europe/Berlin", "de-DE", "5.6.7.8"))
 @patch("cloakbrowser.browser.ensure_binary", return_value="/fake/chrome")
 def test_persistent_context_geoip(_mock_bin, _mock_geoip):