diff --git a/surfsense_backend/app/utils/proxy/__init__.py b/surfsense_backend/app/utils/proxy/__init__.py index 61af110a4..34d33f35d 100644 --- a/surfsense_backend/app/utils/proxy/__init__.py +++ b/surfsense_backend/app/utils/proxy/__init__.py @@ -15,9 +15,16 @@ def get_proxy_url() -> str | None: return get_active_provider().get_proxy_url() -def get_sticky_proxy_url(session_id: str) -> str | None: +def get_geo_proxy_url(country: str | None = None) -> str | None: + """Proxy URL pinned to an exit country when the provider supports it.""" + return get_active_provider().get_geo_proxy_url(country) + + +def get_sticky_proxy_url( + session_id: str, country: str | None = None +) -> str | None: """Proxy URL pinned to a stable vendor session when supported.""" - return get_active_provider().get_sticky_proxy_url(session_id) + return get_active_provider().get_sticky_proxy_url(session_id, country) def get_playwright_proxy() -> dict[str, str] | None: @@ -46,6 +53,7 @@ def get_residential_proxy_url() -> str | None: __all__ = [ "ProxyProvider", "get_active_provider", + "get_geo_proxy_url", "get_playwright_proxy", "get_proxy_url", "get_requests_proxies", diff --git a/surfsense_backend/app/utils/proxy/base.py b/surfsense_backend/app/utils/proxy/base.py index a4ef768c7..f225dcfb1 100644 --- a/surfsense_backend/app/utils/proxy/base.py +++ b/surfsense_backend/app/utils/proxy/base.py @@ -67,13 +67,23 @@ class ProxyProvider(ABC): return None return {"http": proxy_url, "https": proxy_url} - def get_sticky_proxy_url(self, session_id: str) -> str | None: + def get_geo_proxy_url(self, country: str | None = None) -> str | None: + """Return a proxy URL pinned to ``country`` when supported. + + Providers without vendor-specific country routing safely fall back to + their ordinary proxy URL. + """ + return self.get_proxy_url() + + def get_sticky_proxy_url( + self, session_id: str, country: str | None = None + ) -> str | None: """Return a proxy URL pinned to ``session_id`` when supported. Providers without vendor-specific session routing safely fall back to their ordinary proxy URL. """ - return self.get_proxy_url() + return self.get_geo_proxy_url(country) def get_location(self) -> str: """Return the proxy's configured exit region (e.g. ``"us"``), or ``""``. diff --git a/surfsense_backend/app/utils/proxy/providers/dataimpulse.py b/surfsense_backend/app/utils/proxy/providers/dataimpulse.py index ed498321e..b2a151e03 100644 --- a/surfsense_backend/app/utils/proxy/providers/dataimpulse.py +++ b/surfsense_backend/app/utils/proxy/providers/dataimpulse.py @@ -32,9 +32,24 @@ logger = logging.getLogger(__name__) # DataImpulse encodes country routing as a "__cr." username suffix; the # country token runs until the next "__" param (e.g. "__sid") or the end. _COUNTRY_MARKER = "__cr." +_COUNTRY_RE = re.compile(r"__cr\.[A-Za-z]{2,}") _SESSION_RE = re.compile(r"__sid\.[A-Za-z0-9_-]+") +def _safe_session_id(session_id: str) -> str: + safe_id = re.sub(r"[^A-Za-z0-9_-]", "-", session_id).strip("-") + if not safe_id: + raise ValueError("session_id must contain at least one letter or digit") + return safe_id + + +def _safe_country(country: str | None) -> str | None: + if country is None: + return None + safe_country = re.sub(r"[^a-z]", "", country.lower()) + return safe_country or None + + class DataImpulseProvider(ProxyProvider): """Provider for a DataImpulse proxy URL in the shared ``PROXY_URL`` env.""" @@ -54,17 +69,21 @@ class DataImpulseProvider(ProxyProvider): return "" return username.split(_COUNTRY_MARKER, 1)[1].split("__", 1)[0].lower() - def get_sticky_proxy_url(self, session_id: str) -> str | None: - """Return the configured URL with a deterministic sticky-session suffix.""" + def _rewrite_proxy_url( + self, *, country: str | None = None, session_id: str | None = None + ) -> str | None: url = self.get_proxy_url() if not url: return None - safe_id = re.sub(r"[^A-Za-z0-9_-]", "-", session_id).strip("-") - if not safe_id: - raise ValueError("session_id must contain at least one letter or digit") parts = urlsplit(url) - username = parts.username or "" - username = _SESSION_RE.sub("", username) + f"__sid.{safe_id}" + username = _SESSION_RE.sub("", _COUNTRY_RE.sub("", parts.username or "")) + safe_country = _safe_country(country) + if safe_country is not None: + username += f"__cr.{safe_country}" + elif _COUNTRY_MARKER in (parts.username or ""): + username += f"__cr.{self.get_location()}" + if session_id is not None: + username += f"__sid.{_safe_session_id(session_id)}" userinfo = quote(username, safe="%") if parts.password is not None: userinfo += f":{quote(parts.password, safe='%')}" @@ -75,3 +94,15 @@ class DataImpulseProvider(ProxyProvider): return urlunsplit( (parts.scheme, netloc, parts.path, parts.query, parts.fragment) ) + + def get_geo_proxy_url(self, country: str | None = None) -> str | None: + """Return the configured URL with a country-routing suffix when requested.""" + if not _safe_country(country): + return self.get_proxy_url() + return self._rewrite_proxy_url(country=country) + + def get_sticky_proxy_url( + self, session_id: str, country: str | None = None + ) -> str | None: + """Return the configured URL with deterministic country/session suffixes.""" + return self._rewrite_proxy_url(country=country, session_id=session_id)