diff --git a/.env.example b/.env.example index 0aa86dc..44b3953 100644 --- a/.env.example +++ b/.env.example @@ -237,6 +237,13 @@ ZEPTO_MCP_ENABLED=false # API key for production servers (optional for local dev) # SANDBOX_API_KEY= # +# Private Agent Browser takeover. Keep this URL reachable only through +# Tailscale or another authenticated HTTPS reverse proxy. The fragment token +# is sent directly to Telegram and never returned to the model. +# BROWSER_TAKEOVER_PUBLIC_URL=https://blacki.example.ts.net/browser-takeover +# BROWSER_TAKEOVER_TTL_SECONDS=300 +# BROWSER_TAKEOVER_STREAM_PORT=9223 +# # Sandbox TTL in minutes (default: 30) # SANDBOX_TIMEOUT_MINUTES=30 # diff --git a/docs/architecture.md b/docs/architecture.md index 10e0e6a..149dae1 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -216,6 +216,28 @@ An authenticated capability must be implemented as a separately authorized, short-lived broker operation. Adding a standing environment variable is not an acceptable opt-in path. +### Browser authentication takeover + +When `BROWSER_TAKEOVER_PUBLIC_URL` is configured, the private Telegram root +agent can pause an Agent Browser operation and hand the live browser viewport +to the authenticated Telegram user. Blacki sends the one-time link directly +through the Telegram API; the tool result returned to the model never contains +the link or its token. + +The browser page redeems the fragment token once, replaces it with an +HttpOnly, SameSite cookie, and proxies only Agent Browser's viewport and input +protocol. Passwords, OTPs, and typed text travel as transient browser +input events. They are not ADK messages, tool arguments, session events, or log +fields. Completing the takeover closes the stream before the model resumes. + +Takeover state is process-local, hashed, bound to one private Telegram +conversation, and expires after five minutes by default. The upstream sandbox +endpoint and routing headers are never sent to the browser client. The public +URL must sit behind authenticated HTTPS, normally a Tailscale ACL. Browser +cookies and local storage remain sensitive state inside the session sandbox. +Do not enable traces, screenshots, HAR capture, or video recording during a +takeover. + ## Health semantics `/live` is side-effect-free and process-only. `/ready` checks the already diff --git a/docs/base-infra/environment-variables.md b/docs/base-infra/environment-variables.md index 1b512a8..7c68972 100644 --- a/docs/base-infra/environment-variables.md +++ b/docs/base-infra/environment-variables.md @@ -99,6 +99,9 @@ its presence changes model routing. | `CLOUDFLARE_API_TOKEN` | unset | Cloudflare Workers AI API token for Telegram voice transcription | | `KOKORO_TTS_BASE_URL` | unset | Register private Kokoro speech delivery for Telegram | | `KOKORO_TTS_VOICE` | `af_heart` | Kokoro voice ID used for generated MP3 audio | +| `BROWSER_TAKEOVER_PUBLIC_URL` | unset | HTTPS URL for private Agent Browser takeover | +| `BROWSER_TAKEOVER_TTL_SECONDS` | `300` | One-time takeover lifetime, from 60 to 900 seconds | +| `BROWSER_TAKEOVER_STREAM_PORT` | `9223` | Agent Browser stream port inside each sandbox | The token is required and format-validated when Telegram is enabled. When `TELEGRAM_ACCESS_CODE` is set, new users enter it with `/start `; @@ -106,6 +109,12 @@ historical private chats with existing Blacki sessions remain authorized, while groups and topics are rejected. Rotating the code requires code-authorized users to authenticate again without deleting their stored Blacki data. +Browser takeover is enabled only when `BROWSER_TAKEOVER_PUBLIC_URL` is set. +Put that route behind an authenticated HTTPS reverse proxy on Tailscale. The +tool sends its one-time fragment link directly to the private Telegram chat, +waits for the user to return control, and never includes the link in model +context. The sandbox image must provide the `agent-browser` executable. + `KOKORO_TTS_BASE_URL` is an optional HTTP or HTTPS base URL without a path to `/v1/audio/speech`; Blacki appends that fixed endpoint. The URL must be reachable from inside the Blacki container. Do not use `localhost` for a diff --git a/docs/browser-takeover.md b/docs/browser-takeover.md new file mode 100644 index 0000000..e13d52c --- /dev/null +++ b/docs/browser-takeover.md @@ -0,0 +1,66 @@ +# Private browser takeover + +Browser takeover lets the private Telegram user type passwords and one-time +codes into Agent Browser without sending those values through Telegram messages +or ADK model context. + +## How it works + +1. The root agent calls `start_browser_takeover` with an HTTPS login page. +2. Blacki starts Agent Browser's interactive stream in the existing session + sandbox. +3. Blacki sends a single-use link directly to the same private Telegram chat. + The model receives only the eventual completion status. +4. The link opens a live canvas. Mouse, touch, keyboard, and scroll events pass + through Blacki to the Agent Browser stream without entering ADK session + history. +5. The takeover page shows the expected login origin and the current browser + origin outside the streamed page. If navigation changes origin, the page + warns the user to verify the new site before typing sensitive data. +6. Clicking **Done** disables the stream and resumes the waiting tool call. + Disconnects, expiry, replacement, and process shutdown do not count as + successful human completion. + +The URL fragment is exchanged once for an HttpOnly, SameSite cookie. Blacki +stores only token hashes, sandbox routing details, the expected origin, and +completion state in process memory. Restarting Blacki invalidates every +takeover. + +## Configure + +The sandbox image must contain a compatible `agent-browser` executable. Set: + +```dotenv +BROWSER_TAKEOVER_PUBLIC_URL=https://blacki.example.ts.net/browser-takeover +BROWSER_TAKEOVER_TTL_SECONDS=300 +BROWSER_TAKEOVER_STREAM_PORT=9223 +``` + +The public URL must end with `/browser-takeover`. HTTPS is mandatory except on +loopback during development. Keep the endpoint on your tailnet and restrict it +with Tailscale ACLs or equivalent device authentication. The reverse proxy must +forward WebSocket upgrades on `/browser-takeover/ws`. + +Blacki's production Compose override publishes the application on +`127.0.0.1` by default. Point Caddy or Tailscale Serve at that loopback port; +do not publish the application port on every host interface. + +## Operational limits + +- One takeover can run per Telegram conversation. +- Links are single-use and expire after five minutes by default. +- Only matching private-chat sender and chat identities can create a takeover. +- A model turn remains paused until the user clicks **Done** or the link expires. +- Starting a new Telegram turn may cancel the waiting model turn. +- Takeover capability state is process-local. Run one application worker and + one replica, or guarantee sticky routing for the originating tool execution, + `/redeem`, `/ws`, and `/complete`. A multi-worker or multi-replica deployment + without shared takeover state will intermittently reject valid links. +- Cookies and local storage created by the target website remain inside the + session sandbox and are as sensitive as the original password. + +Do not enable screenshots, video, HAR files, browser command tracing, or debug +payload logging while a takeover is active. Blacki does not record the input +events, but the target website and Telegram still have their ordinary service +metadata. This feature keeps credentials out of LLM and Blacki content logs; it +does not claim that no system involved in the login has operational logs. diff --git a/mkdocs.yml b/mkdocs.yml index 2e14d15..e444bfc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -51,6 +51,7 @@ nav: - Telegram setup: telegram-setup.md - Zepto MCP: zepto-mcp.md - Gmail API: gmail.md + - Private browser takeover: browser-takeover.md - Local development: development.md - Operate: - Configuration: base-infra/environment-variables.md diff --git a/pyproject.toml b/pyproject.toml index 652ca63..ca9a3b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "mem0ai>=0.1.0,<1.0.0", "opensandbox-code-interpreter>=0.1.2", "boto3>=1.40.0,<2.0.0", + "websockets>=15.0.1,<16.0.0", ] [project.scripts] diff --git a/src/blacki/browser_takeover/__init__.py b/src/blacki/browser_takeover/__init__.py new file mode 100644 index 0000000..1533fa1 --- /dev/null +++ b/src/blacki/browser_takeover/__init__.py @@ -0,0 +1,20 @@ +"""Private human takeover for Agent Browser sessions.""" + +from .config import BrowserTakeoverConfig, BrowserTakeoverConfigurationError +from .service import ( + BrowserTakeoverError, + BrowserTakeoverLease, + BrowserTakeoverService, + get_browser_takeover_service, + reset_browser_takeover_service, +) + +__all__ = [ + "BrowserTakeoverConfig", + "BrowserTakeoverConfigurationError", + "BrowserTakeoverError", + "BrowserTakeoverLease", + "BrowserTakeoverService", + "get_browser_takeover_service", + "reset_browser_takeover_service", +] diff --git a/src/blacki/browser_takeover/config.py b/src/blacki/browser_takeover/config.py new file mode 100644 index 0000000..9bc57ca --- /dev/null +++ b/src/blacki/browser_takeover/config.py @@ -0,0 +1,87 @@ +"""Configuration for the private browser takeover endpoint.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from urllib.parse import urlsplit + + +class BrowserTakeoverConfigurationError(ValueError): + """Raised when browser takeover configuration is unsafe or incomplete.""" + + +@dataclass(frozen=True) +class BrowserTakeoverConfig: + """Validated browser takeover configuration.""" + + public_url: str + ttl_seconds: int = 300 + stream_port: int = 9223 + + @classmethod + def from_environment(cls) -> BrowserTakeoverConfig | None: + """Build configuration when the optional takeover URL is present.""" + public_url = os.getenv("BROWSER_TAKEOVER_PUBLIC_URL", "").strip() + if not public_url: + return None + + try: + ttl_seconds = int(os.getenv("BROWSER_TAKEOVER_TTL_SECONDS", "300")) + stream_port = int(os.getenv("BROWSER_TAKEOVER_STREAM_PORT", "9223")) + except ValueError as exc: + raise BrowserTakeoverConfigurationError( + "Browser takeover TTL and stream port must be integers" + ) from exc + + config = cls( + public_url=public_url.rstrip("/"), + ttl_seconds=ttl_seconds, + stream_port=stream_port, + ) + config.validate() + return config + + def validate(self) -> None: + """Reject public URLs and limits that could expose a browser session.""" + parsed = urlsplit(self.public_url) + is_loopback_http = parsed.scheme == "http" and parsed.hostname in { + "127.0.0.1", + "localhost", + } + if parsed.scheme != "https" and not is_loopback_http: + raise BrowserTakeoverConfigurationError( + "BROWSER_TAKEOVER_PUBLIC_URL must use HTTPS" + ) + if not parsed.netloc or parsed.username or parsed.password: + raise BrowserTakeoverConfigurationError( + "BROWSER_TAKEOVER_PUBLIC_URL must be an absolute URL " + "without credentials" + ) + if parsed.query or parsed.fragment: + raise BrowserTakeoverConfigurationError( + "BROWSER_TAKEOVER_PUBLIC_URL cannot contain a query or fragment" + ) + if parsed.path != "/browser-takeover": + raise BrowserTakeoverConfigurationError( + "BROWSER_TAKEOVER_PUBLIC_URL must end with /browser-takeover" + ) + if not 60 <= self.ttl_seconds <= 900: + raise BrowserTakeoverConfigurationError( + "BROWSER_TAKEOVER_TTL_SECONDS must be between 60 and 900" + ) + if not 1024 <= self.stream_port <= 65535: + raise BrowserTakeoverConfigurationError( + "BROWSER_TAKEOVER_STREAM_PORT must be between 1024 and 65535" + ) + + @property + def public_origin(self) -> str: + """Return the exact origin allowed to redeem takeover tokens.""" + parsed = urlsplit(self.public_url) + return f"{parsed.scheme}://{parsed.netloc}" + + @property + def secure_cookie(self) -> bool: + """Require Secure cookies outside loopback development.""" + return urlsplit(self.public_url).scheme == "https" diff --git a/src/blacki/browser_takeover/routes.py b/src/blacki/browser_takeover/routes.py new file mode 100644 index 0000000..e9168c1 --- /dev/null +++ b/src/blacki/browser_takeover/routes.py @@ -0,0 +1,647 @@ +"""FastAPI routes for private Agent Browser pair browsing.""" + +from __future__ import annotations + +import asyncio +import json +import math +from typing import Any, TypeGuard +from urllib.parse import urlsplit + +from fastapi import APIRouter, Request, WebSocket, WebSocketDisconnect +from fastapi.responses import HTMLResponse, JSONResponse, Response +from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosed + +from .service import BrowserTakeoverService, get_browser_takeover_service + +COOKIE_NAME = "blacki_browser_takeover" +_MAX_INPUT_BYTES = 16 * 1024 +_MAX_TEXT_BYTES = 4096 +_ALLOWED_MOUSE_EVENTS = frozenset( + {"mousePressed", "mouseReleased", "mouseMoved", "mouseWheel"} +) +_ALLOWED_MOUSE_BUTTONS = frozenset({"left", "middle", "right", "none"}) +_ALLOWED_KEYBOARD_EVENTS = frozenset({"keyDown", "keyUp", "char"}) +_ALLOWED_TOUCH_EVENTS = frozenset({"touchStart", "touchMove", "touchEnd"}) +_FRAME_METADATA_KEYS = frozenset( + { + "deviceWidth", + "deviceHeight", + "pageScaleFactor", + "offsetTop", + "scrollOffsetX", + "scrollOffsetY", + "timestamp", + } +) +_SECURITY_HEADERS = { + "Cache-Control": "no-store, max-age=0", + "Content-Security-Policy": ( + "default-src 'none'; connect-src 'self' ws: wss:; img-src data:; " + "script-src 'unsafe-inline'; style-src 'unsafe-inline'; " + "base-uri 'none'; form-action 'none'; frame-ancestors 'none'" + ), + "Permissions-Policy": "camera=(), microphone=(), geolocation=()", + "Referrer-Policy": "no-referrer", + "X-Content-Type-Options": "nosniff", + "X-Frame-Options": "DENY", +} + +_PAGE = """ + + + + + Private browser control + + + +
+
+ Private browser control + Expected site: verifying... + Current site: waiting for browser... +
+ Connecting... +
+ + + + +""" + + +def _service() -> BrowserTakeoverService | None: + return get_browser_takeover_service() + + +def _same_origin(request: Request, service: BrowserTakeoverService) -> bool: + return request.headers.get("origin") == service.config.public_origin + + +def _private_response(content: Any, *, status_code: int = 200) -> JSONResponse: + return JSONResponse(content, status_code=status_code, headers=_SECURITY_HEADERS) + + +def _is_int(value: Any) -> TypeGuard[int]: + return isinstance(value, int) and not isinstance(value, bool) + + +def _is_number(value: Any) -> TypeGuard[int | float]: + return ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + ) + + +def _bounded_number(value: Any, *, limit: float = 100_000.0) -> bool: + return _is_number(value) and -limit <= value <= limit + + +def _only_keys(payload: dict[str, Any], allowed: frozenset[str]) -> bool: + return all(key in allowed for key in payload) + + +def _valid_ack(payload: dict[str, Any]) -> bool: + if not _only_keys(payload, frozenset({"type", "seq"})): + return False + seq = payload.get("seq") + return _is_int(seq) and 0 <= seq <= 2**53 - 1 + + +def _valid_config(payload: dict[str, Any]) -> bool: + if not _only_keys(payload, frozenset({"type", "maxFps", "pacing"})): + return False + if "maxFps" not in payload and "pacing" not in payload: + return False + max_fps = payload.get("maxFps") + if max_fps is not None and (not _is_int(max_fps) or not 0 <= max_fps <= 120): + return False + pacing = payload.get("pacing") + return pacing is None or pacing == "ack" + + +def _short_string(value: Any, *, max_bytes: int = 128) -> bool: + return isinstance(value, str) and len(value.encode()) <= max_bytes + + +def _valid_keyboard(payload: dict[str, Any]) -> bool: + allowed = frozenset({"type", "eventType", "key", "code", "text", "modifiers"}) + if not _only_keys(payload, allowed): + return False + event_type = payload.get("eventType") + if event_type not in _ALLOWED_KEYBOARD_EVENTS: + return False + modifiers = payload.get("modifiers") + if modifiers is not None and (not _is_int(modifiers) or not 0 <= modifiers <= 15): + return False + if event_type == "char": + text = payload.get("text") + return ( + isinstance(text, str) + and bool(text) + and len(text.encode()) <= _MAX_TEXT_BYTES + ) + key = payload.get("key") + code = payload.get("code") + return _short_string(key) and bool(key) and (code is None or _short_string(code)) + + +def _valid_mouse(payload: dict[str, Any]) -> bool: + allowed = frozenset( + { + "type", + "eventType", + "x", + "y", + "button", + "clickCount", + "deltaX", + "deltaY", + "modifiers", + } + ) + if not _only_keys(payload, allowed): + return False + event_type = payload.get("eventType") + if event_type not in _ALLOWED_MOUSE_EVENTS: + return False + if not _bounded_number(payload.get("x")) or not _bounded_number(payload.get("y")): + return False + modifiers = payload.get("modifiers") + if modifiers is not None and (not _is_int(modifiers) or not 0 <= modifiers <= 15): + return False + button = payload.get("button") + if button is not None and button not in _ALLOWED_MOUSE_BUTTONS: + return False + click_count = payload.get("clickCount") + if click_count is not None and ( + not _is_int(click_count) or not 0 <= click_count <= 3 + ): + return False + if event_type == "mouseWheel": + if "deltaX" not in payload and "deltaY" not in payload: + return False + return ("deltaX" not in payload or _bounded_number(payload.get("deltaX"))) and ( + "deltaY" not in payload or _bounded_number(payload.get("deltaY")) + ) + return True + + +def _valid_touch_point(point: Any) -> bool: + if not isinstance(point, dict): + return False + if not _only_keys(point, frozenset({"x", "y", "id"})): + return False + if not _bounded_number(point.get("x")) or not _bounded_number(point.get("y")): + return False + touch_id = point.get("id") + return touch_id is None or (_is_int(touch_id) and 0 <= touch_id <= 2**31 - 1) + + +def _valid_touch(payload: dict[str, Any]) -> bool: + if not _only_keys(payload, frozenset({"type", "eventType", "touchPoints"})): + return False + event_type = payload.get("eventType") + if event_type not in _ALLOWED_TOUCH_EVENTS: + return False + points = payload.get("touchPoints") + if not isinstance(points, list) or len(points) > 10: + return False + if event_type != "touchEnd" and not points: + return False + return all(_valid_touch_point(point) for point in points) + + +def _valid_client_message(message: str) -> bool: + if len(message.encode()) > _MAX_INPUT_BYTES: + return False + try: + payload = json.loads(message) + except json.JSONDecodeError: + return False + if not isinstance(payload, dict): + return False + message_type = payload.get("type") + if message_type == "ack": + return _valid_ack(payload) + if message_type == "config": + return _valid_config(payload) + if message_type == "input_keyboard": + return _valid_keyboard(payload) + if message_type == "input_mouse": + return _valid_mouse(payload) + if message_type == "input_touch": + return _valid_touch(payload) + return False + + +def _normalized_web_origin(url: str) -> str | None: + try: + parsed = urlsplit(url) + host = parsed.hostname + port = parsed.port + except ValueError: + return None + if parsed.scheme not in {"http", "https"} or host is None: + return None + host = host.lower() + if ":" in host: + host = f"[{host}]" + default_port = (parsed.scheme == "https" and port == 443) or ( + parsed.scheme == "http" and port == 80 + ) + suffix = "" if port is None or default_port else f":{port}" + return f"{parsed.scheme.lower()}://{host}{suffix}" + + +def _safe_server_message(message: str) -> str | None: + try: + payload = json.loads(message) + except json.JSONDecodeError: + return None + if not isinstance(payload, dict): + return None + message_type = payload.get("type") + if message_type == "frame": + data = payload.get("data") + if not isinstance(data, str): + return None + safe: dict[str, Any] = {"type": "frame", "data": data} + seq = payload.get("seq") + if _is_int(seq) and seq >= 0: + safe["seq"] = seq + metadata = payload.get("metadata") + if isinstance(metadata, dict): + safe_metadata = { + key: value + for key, value in metadata.items() + if key in _FRAME_METADATA_KEYS and _is_number(value) + } + if safe_metadata: + safe["metadata"] = safe_metadata + return json.dumps(safe, separators=(",", ":")) + if message_type == "status": + connected = payload.get("connected") + if not isinstance(connected, bool): + return None + return json.dumps( + {"type": "status", "connected": connected}, + separators=(",", ":"), + ) + if message_type == "url": + url = payload.get("url") + origin = _normalized_web_origin(url) if isinstance(url, str) else None + if origin is None: + return None + return json.dumps( + {"type": "navigation", "origin": origin}, + separators=(",", ":"), + ) + return None + + +def create_browser_takeover_router() -> APIRouter: + """Create routes without exposing upstream sandbox endpoints to the client.""" + router = APIRouter(prefix="/browser-takeover") + + @router.get("", response_class=HTMLResponse) + async def takeover_page() -> HTMLResponse: + if _service() is None: + return HTMLResponse("Not found", status_code=404) + return HTMLResponse(_PAGE, headers=_SECURITY_HEADERS) + + @router.post("/redeem") + async def redeem(request: Request) -> Response: + service = _service() + if service is None or not _same_origin(request, service): + return _private_response({"status": "invalid"}, status_code=404) + try: + raw_body = await request.body() + if len(raw_body) > 4096: + return _private_response({"status": "invalid"}, status_code=404) + body = json.loads(raw_body) + token = body.get("token", "") if isinstance(body, dict) else "" + if not isinstance(token, str): + token = "" + except (json.JSONDecodeError, UnicodeDecodeError): + token = "" + browser_token = await service.redeem(token) + if browser_token is None: + return _private_response({"status": "invalid"}, status_code=404) + response = _private_response({"status": "ready"}) + response.set_cookie( + COOKIE_NAME, + browser_token, + httponly=True, + secure=service.config.secure_cookie, + samesite="strict", + max_age=service.config.ttl_seconds, + path="/browser-takeover", + ) + return response + + @router.post("/complete") + async def complete(request: Request) -> Response: + service = _service() + token = request.cookies.get(COOKIE_NAME) + if ( + service is None + or not _same_origin(request, service) + or not await service.complete(token) + ): + return _private_response({"status": "invalid"}, status_code=404) + response = _private_response({"status": "complete"}) + response.delete_cookie(COOKIE_NAME, path="/browser-takeover") + return response + + @router.websocket("/ws") + async def browser_stream(websocket: WebSocket) -> None: + service = _service() + token = websocket.cookies.get(COOKIE_NAME) + same_origin = ( + service is not None + and websocket.headers.get("origin") == service.config.public_origin + ) + if service is None or not same_origin: + await websocket.close(code=4401) + return + session = await service.authorize(token) + if session is None: + await websocket.close(code=4401) + return + + await websocket.accept() + await websocket.send_json( + { + "type": "takeover_meta", + "expectedOrigin": session.expected_origin, + } + ) + try: + async with connect( + session.upstream_url, + additional_headers=session.upstream_headers, + max_size=8 * 1024 * 1024, + ) as upstream: + + async def from_browser() -> None: + async for message in upstream: + if isinstance(message, bytes): + await websocket.send_bytes(message) + continue + safe_message = _safe_server_message(message) + if safe_message is not None: + await websocket.send_text(safe_message) + + async def from_user() -> None: + while True: + message = await websocket.receive_text() + if _valid_client_message(message): + await upstream.send(message) + + tasks = { + asyncio.create_task(from_browser()), + asyncio.create_task(from_user()), + } + done, pending = await asyncio.wait( + tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + await asyncio.gather(*done, *pending, return_exceptions=True) + except (ConnectionClosed, WebSocketDisconnect): + return + except Exception: + await websocket.close(code=1011) + + return router diff --git a/src/blacki/browser_takeover/service.py b/src/blacki/browser_takeover/service.py new file mode 100644 index 0000000..f0dd551 --- /dev/null +++ b/src/blacki/browser_takeover/service.py @@ -0,0 +1,351 @@ +"""In-memory, user-bound browser takeover lifecycle.""" + +from __future__ import annotations + +import asyncio +import hashlib +import secrets +import time +from dataclasses import dataclass, field +from datetime import timedelta +from typing import Any +from urllib.parse import urlsplit + +from opensandbox.models.execd import RunCommandOpts + +from blacki.sandbox.manager import get_sandbox_manager + +from .config import BrowserTakeoverConfig + +_START_COMMAND = ( + 'agent-browser open "$BLACKI_BROWSER_TAKEOVER_URL" ' + '&& agent-browser stream enable --port "$BLACKI_BROWSER_STREAM_PORT"' +) +_STOP_COMMAND = "agent-browser stream disable" + + +class BrowserTakeoverError(RuntimeError): + """A secret-free browser takeover failure.""" + + +@dataclass +class _TakeoverSession: + session_id: str + owner_key: str + sandbox: Any = field(repr=False) + upstream_url: str = field(repr=False) + upstream_headers: dict[str, str] = field(repr=False) + expected_origin: str + expires_at: float + finished: asyncio.Event = field(default_factory=asyncio.Event, repr=False) + completed: bool = False + browser_token_digest: str | None = field(default=None, repr=False) + + +@dataclass(frozen=True) +class BrowserTakeoverLease: + """Internal lease returned to the tool without exposing tokens to the model.""" + + session_id: str + takeover_url: str = field(repr=False) + expires_in_seconds: int + + +def _digest(value: str) -> str: + return hashlib.sha256(value.encode()).hexdigest() + + +def _websocket_url(endpoint: str) -> str: + if "://" not in endpoint: + return f"ws://{endpoint}" + parsed = urlsplit(endpoint) + scheme = "wss" if parsed.scheme in {"https", "wss"} else "ws" + return parsed._replace(scheme=scheme).geturl() + + +def _validated_login_url(login_url: str) -> str: + parsed = urlsplit(login_url.strip()) + if parsed.scheme != "https" or not parsed.netloc: + raise BrowserTakeoverError("Browser takeover requires an HTTPS login URL") + if parsed.username or parsed.password: + raise BrowserTakeoverError("Login URLs cannot contain credentials") + return parsed.geturl() + + +def _origin(url: str) -> str: + parsed = urlsplit(url) + host = parsed.hostname + if host is None: + raise BrowserTakeoverError("Browser takeover requires a valid login host") + try: + port = parsed.port + except ValueError as exc: + raise BrowserTakeoverError( + "Browser takeover login URL has an invalid port" + ) from exc + host = host.lower() + if ":" in host: + host = f"[{host}]" + default_port = (parsed.scheme == "https" and port == 443) or ( + parsed.scheme == "http" and port == 80 + ) + port_suffix = "" if port is None or default_port else f":{port}" + return f"{parsed.scheme.lower()}://{host}{port_suffix}" + + +class BrowserTakeoverService: + """Create one-time links and proxy only the matching browser stream.""" + + def __init__(self, config: BrowserTakeoverConfig) -> None: + config.validate() + self.config = config + self._sessions: dict[str, _TakeoverSession] = {} + self._takeover_tokens: dict[str, str] = {} + self._browser_tokens: dict[str, str] = {} + self._owners: dict[str, str] = {} + self._starting_owners: set[str] = set() + self._lock = asyncio.Lock() + + async def create(self, *, login_url: str, state: Any) -> BrowserTakeoverLease: + """Start streaming the session browser and mint a single-use link.""" + safe_url = _validated_login_url(login_url) + expected_origin = _origin(safe_url) + owner_key = self._owner_key(state) + expired_session: _TakeoverSession | None = None + async with self._lock: + active_id = self._owners.get(owner_key) + active = self._sessions.get(active_id or "") + if active is not None and not self._expired(active): + raise BrowserTakeoverError( + "A browser takeover is already active for this conversation" + ) + if active is not None: + self._remove_session_locked(active) + expired_session = active + if owner_key in self._starting_owners: + raise BrowserTakeoverError( + "A browser takeover is already starting for this conversation" + ) + self._starting_owners.add(owner_key) + + sandbox: Any | None = None + try: + if expired_session is not None: + await self._stop_stream(expired_session.sandbox) + manager = get_sandbox_manager() + result = await manager.get_or_create_sandbox(state) + sandbox = result.get("sandbox") + if sandbox is None: + raise BrowserTakeoverError("The browser sandbox is unavailable") + + opts = RunCommandOpts( + timeout=timedelta(seconds=60), + envs={ + "BLACKI_BROWSER_TAKEOVER_URL": safe_url, + "BLACKI_BROWSER_STREAM_PORT": str(self.config.stream_port), + }, + ) + execution = await sandbox.commands.run(_START_COMMAND, opts=opts) + if execution.error: + raise BrowserTakeoverError("Agent Browser takeover could not start") + endpoint = await sandbox.get_endpoint(self.config.stream_port) + except BrowserTakeoverError: + await self._abort_start(owner_key, sandbox) + raise + except asyncio.CancelledError: + await self._abort_start(owner_key, sandbox) + raise + except Exception as exc: + await self._abort_start(owner_key, sandbox) + raise BrowserTakeoverError( + "Agent Browser takeover could not start" + ) from exc + + session_id = secrets.token_urlsafe(24) + takeover_token = secrets.token_urlsafe(32) + session = _TakeoverSession( + session_id=session_id, + owner_key=owner_key, + sandbox=sandbox, + upstream_url=_websocket_url(endpoint.endpoint), + upstream_headers=dict(endpoint.headers), + expected_origin=expected_origin, + expires_at=time.monotonic() + self.config.ttl_seconds, + ) + async with self._lock: + self._starting_owners.discard(owner_key) + self._sessions[session_id] = session + self._takeover_tokens[_digest(takeover_token)] = session_id + self._owners[owner_key] = session_id + + return BrowserTakeoverLease( + session_id=session_id, + takeover_url=f"{self.config.public_url}#{takeover_token}", + expires_in_seconds=self.config.ttl_seconds, + ) + + async def redeem(self, takeover_token: str) -> str | None: + """Consume a link token and return a cookie token exactly once.""" + if not isinstance(takeover_token, str) or not takeover_token: + return None + if len(takeover_token) > 256: + return None + expired_session: _TakeoverSession | None = None + async with self._lock: + session_id = self._takeover_tokens.pop(_digest(takeover_token), None) + session = self._sessions.get(session_id or "") + if session is None or self._expired(session): + if session is not None: + self._remove_session_locked(session) + expired_session = session + else: + browser_token = secrets.token_urlsafe(32) + digest = _digest(browser_token) + session.browser_token_digest = digest + self._browser_tokens[digest] = session.session_id + return browser_token + if expired_session is not None: + await self._stop_stream(expired_session.sandbox) + return None + + async def authorize(self, browser_token: str | None) -> _TakeoverSession | None: + """Resolve an active browser cookie without exposing its value.""" + if not browser_token: + return None + expired_session: _TakeoverSession | None = None + async with self._lock: + session_id = self._browser_tokens.get(_digest(browser_token)) + session = self._sessions.get(session_id or "") + if session is None or self._expired(session): + if session is not None: + self._remove_session_locked(session) + expired_session = session + else: + return session + if expired_session is not None: + await self._stop_stream(expired_session.sandbox) + return None + + async def complete(self, browser_token: str | None) -> bool: + """End human control and wake the waiting agent tool.""" + session = await self.authorize(browser_token) + if session is None: + return False + session.completed = True + session.finished.set() + return True + + async def wait(self, lease: BrowserTakeoverLease) -> bool: + """Wait for explicit user completion or for the lease to end.""" + async with self._lock: + session = self._sessions.get(lease.session_id) + if session is None: + return False + remaining = max(0.0, session.expires_at - time.monotonic()) + try: + await asyncio.wait_for(session.finished.wait(), timeout=remaining) + return session.completed + except TimeoutError: + return False + + async def close(self, lease: BrowserTakeoverLease) -> None: + """Disable streaming and erase every in-memory capability.""" + async with self._lock: + session = self._sessions.get(lease.session_id) + if session is not None: + self._remove_session_locked(session) + if session is None: + return + await self._stop_stream(session.sandbox) + + async def close_all(self) -> None: + """Invalidate all links during application shutdown.""" + async with self._lock: + sessions = list(self._sessions.values()) + self._sessions.clear() + self._takeover_tokens.clear() + self._browser_tokens.clear() + self._owners.clear() + self._starting_owners.clear() + for session in sessions: + session.finished.set() + stopped: set[int] = set() + for session in sessions: + sandbox_key = id(session.sandbox) + if sandbox_key in stopped: + continue + stopped.add(sandbox_key) + await self._stop_stream(session.sandbox) + + async def _abort_start(self, owner_key: str, sandbox: Any | None) -> None: + if sandbox is not None: + await self._stop_stream(sandbox) + async with self._lock: + self._starting_owners.discard(owner_key) + + @staticmethod + async def _stop_stream(sandbox: Any) -> None: + try: + await sandbox.commands.run( + _STOP_COMMAND, + opts=RunCommandOpts(timeout=timedelta(seconds=10)), + ) + except Exception: # noqa: S112 - cleanup stays silent and secret-free + return + + @staticmethod + def _owner_key(state: Any) -> str: + chat_type = state.get("telegram_chat_type") + chat_id = state.get("telegram_chat_id") + sender_id = state.get("temp:telegram_sender_user_id") + if chat_type != "private" or not chat_id or not sender_id: + raise BrowserTakeoverError( + "Browser takeover is available only in an authenticated " + "private Telegram chat" + ) + if str(chat_id) != str(sender_id): + raise BrowserTakeoverError("Telegram browser takeover identity mismatch") + thread_id = state.get("telegram_thread_id", "") + return f"{chat_id}:{thread_id}" + + @staticmethod + def _expired(session: _TakeoverSession) -> bool: + return time.monotonic() >= session.expires_at + + def _remove_session_locked(self, session: _TakeoverSession) -> None: + self._sessions.pop(session.session_id, None) + self._owners.pop(session.owner_key, None) + stale_takeover = [ + token + for token, session_id in self._takeover_tokens.items() + if session_id == session.session_id + ] + for token in stale_takeover: + self._takeover_tokens.pop(token, None) + if session.browser_token_digest: + self._browser_tokens.pop(session.browser_token_digest, None) + session.finished.set() + + +_service: BrowserTakeoverService | None = None + + +def get_browser_takeover_service() -> BrowserTakeoverService | None: + """Return the configured process-wide takeover service.""" + global _service + if _service is not None: + return _service + config = BrowserTakeoverConfig.from_environment() + if config is None: + return None + _service = BrowserTakeoverService(config) + return _service + + +async def reset_browser_takeover_service() -> None: + """Close and clear the process-wide service.""" + global _service + service = _service + _service = None + if service is not None: + await service.close_all() diff --git a/src/blacki/security/tool_privacy.py b/src/blacki/security/tool_privacy.py index 1d9e465..5409d00 100644 --- a/src/blacki/security/tool_privacy.py +++ b/src/blacki/security/tool_privacy.py @@ -24,6 +24,7 @@ "list_user_files", "restore_user_file", "delete_user_file", + "start_browser_takeover", } ) @@ -70,6 +71,7 @@ def private_tool_privacy_enabled() -> bool: or kokoro_tts_enabled() or google_health_enabled() or r2_files_enabled() + or bool(os.getenv("BROWSER_TAKEOVER_PUBLIC_URL", "").strip()) ) diff --git a/src/blacki/server.py b/src/blacki/server.py index ac86273..1805aaa 100644 --- a/src/blacki/server.py +++ b/src/blacki/server.py @@ -18,6 +18,7 @@ from google.adk.cli.fast_api import get_fast_api_app from openinference.instrumentation.google_adk import GoogleADKInstrumentor +from .browser_takeover.routes import create_browser_takeover_router from .config.paths import agent_root from .container import AppContainer, close_container, init_container from .dashboard.routes import create_dashboard_router @@ -381,6 +382,10 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: from .telegram.progress_callbacks import close_shared_notify_client await close_shared_notify_client() + + from .browser_takeover import reset_browser_takeover_service + + await reset_browser_takeover_service() _shutdown_tracer_provider() @@ -396,6 +401,7 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: ) app.include_router(create_dashboard_router(env)) +app.include_router(create_browser_takeover_router()) @app.get( diff --git a/src/blacki/skills/agent_browser/SKILL.md b/src/blacki/skills/agent_browser/SKILL.md index 72a40d8..d4e0f66 100644 --- a/src/blacki/skills/agent_browser/SKILL.md +++ b/src/blacki/skills/agent_browser/SKILL.md @@ -194,6 +194,12 @@ agent-browser auth save my-app --url https://app.example.com/login \ agent-browser auth login my-app # fills + clicks, waits for form ``` +In a private Telegram session, prefer `start_browser_takeover` when it is +available. It pauses the tool call, sends a one-time interactive browser link +directly to the user, and resumes only after the user returns control. Never ask +the user to send a password or OTP through Telegram. Never put credentials in a +`sandbox_run_command` argument, environment variable, file, or tool result. + ### Persist session across runs ```bash diff --git a/src/blacki/tools/browser_takeover.py b/src/blacki/tools/browser_takeover.py new file mode 100644 index 0000000..fd81c1d --- /dev/null +++ b/src/blacki/tools/browser_takeover.py @@ -0,0 +1,95 @@ +"""Private Telegram tool for handing Agent Browser control to the user.""" + +from __future__ import annotations + +import os +from typing import Any + +from google.adk.tools import ToolContext + +from blacki.browser_takeover import BrowserTakeoverError, get_browser_takeover_service +from blacki.telegram.api import TelegramApiClient +from blacki.telegram.types import InlineKeyboardButton, InlineKeyboardMarkup + + +async def start_browser_takeover( + login_url: str, + tool_context: ToolContext, +) -> dict[str, Any]: + """Let the Telegram user privately complete login in the session browser. + + Use this only when a website requires a password, OTP, CAPTCHA, or other + sensitive human input. The takeover link is sent directly through + Telegram and is deliberately omitted from this tool's result. + """ + service = get_browser_takeover_service() + if service is None: + return { + "status": "unavailable", + "message": "Private browser takeover is not configured.", + } + + try: + lease = await service.create(login_url=login_url, state=tool_context.state) + except BrowserTakeoverError as exc: + return {"status": "error", "message": str(exc)} + + try: + await _send_takeover_link( + lease.takeover_url, + lease.expires_in_seconds, + tool_context.state, + ) + completed = await service.wait(lease) + if not completed: + return { + "status": "expired", + "message": "The private browser takeover expired before completion.", + } + return { + "status": "success", + "message": "The user returned control of the authenticated browser.", + } + except Exception: + return { + "status": "error", + "message": "The private browser takeover could not be delivered.", + } + finally: + await service.close(lease) + + +async def _send_takeover_link( + takeover_url: str, + expires_in_seconds: int, + state: Any, +) -> None: + token = os.getenv("TELEGRAM_BOT_TOKEN", "").strip() + if not token: + raise BrowserTakeoverError("Telegram is not configured") + chat_id = int(state["telegram_chat_id"]) + thread_value = state.get("telegram_thread_id") + thread_id = int(thread_value) if thread_value else None + minutes = max(1, expires_in_seconds // 60) + text = ( + "Private browser control is ready. Open this link, complete the login, " + f"then tap Done. It expires in {minutes} minutes." + ) + async with TelegramApiClient(token) as api: + await api.send_message( + chat_id=chat_id, + text=text, + message_thread_id=thread_id, + protect_content=True, + reply_markup=InlineKeyboardMarkup( + inline_keyboard=[ + [ + InlineKeyboardButton( + text="Take over browser", + callback_data=None, + url=takeover_url, + ) + ] + ] + ), + ) diff --git a/src/blacki/tools/registry.py b/src/blacki/tools/registry.py index b20c207..76abc60 100644 --- a/src/blacki/tools/registry.py +++ b/src/blacki/tools/registry.py @@ -52,6 +52,7 @@ class ToolConfig: zepto_mcp_allowed_chat_ids: frozenset[str] = frozenset() gmail_config: GmailConfig | None = field(default=None, repr=False) r2_files_enabled: bool = False + browser_takeover_enabled: bool = False def build_tools( @@ -113,6 +114,13 @@ def build_tools( if include_user_scoped_tools and config.r2_files_enabled: tools.extend(_build_user_file_tools()) + if ( + include_user_scoped_tools + and config.sandbox_enabled + and config.browser_takeover_enabled + ): + tools.extend(_build_browser_takeover_tools()) + tools.extend(_build_memory_tools()) return tools @@ -371,6 +379,18 @@ def _build_user_file_tools() -> list[Any]: return [] +def _build_browser_takeover_tools() -> list[Any]: + """Build the private Telegram Agent Browser takeover tool.""" + try: + from blacki.tools.browser_takeover import start_browser_takeover + + logger.info("Private browser takeover tool enabled for the Telegram root agent") + return [start_browser_takeover] + except ImportError as exc: # pragma: no cover + logger.warning("Private browser takeover tool disabled: %s", exc) + return [] + + def _build_health_tools() -> list[Any]: """Build the private, read-only Google Health tool.""" try: @@ -432,6 +452,9 @@ def build_tool_config_from_env() -> ToolConfig: gmail_config = GmailConfig.from_environment() except GmailConfigurationError as exc: logger.warning("Gmail API disabled: %s", exc) + from blacki.browser_takeover import BrowserTakeoverConfig + + browser_takeover_enabled = BrowserTakeoverConfig.from_environment() is not None return ToolConfig( exa_api_key=os.getenv("EXA_API_KEY", "").strip() or None, @@ -460,6 +483,7 @@ def build_tool_config_from_env() -> ToolConfig: gmail_config=gmail_config, r2_files_enabled=os.getenv("R2_FILES_ENABLED", "false").strip().lower() in ("true", "1", "yes"), + browser_takeover_enabled=browser_takeover_enabled, ) diff --git a/tests/browser_takeover/__init__.py b/tests/browser_takeover/__init__.py new file mode 100644 index 0000000..205602e --- /dev/null +++ b/tests/browser_takeover/__init__.py @@ -0,0 +1 @@ +"""Browser takeover tests.""" diff --git a/tests/browser_takeover/test_config.py b/tests/browser_takeover/test_config.py new file mode 100644 index 0000000..c83aec5 --- /dev/null +++ b/tests/browser_takeover/test_config.py @@ -0,0 +1,87 @@ +"""Tests for browser takeover configuration.""" + +import pytest + +from blacki.browser_takeover.config import ( + BrowserTakeoverConfig, + BrowserTakeoverConfigurationError, +) + + +def test_config_is_optional(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("BROWSER_TAKEOVER_PUBLIC_URL", raising=False) + + assert BrowserTakeoverConfig.from_environment() is None + + +def test_valid_https_config(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv( + "BROWSER_TAKEOVER_PUBLIC_URL", + "https://blacki.example.ts.net/browser-takeover/", + ) + monkeypatch.setenv("BROWSER_TAKEOVER_TTL_SECONDS", "420") + monkeypatch.setenv("BROWSER_TAKEOVER_STREAM_PORT", "9323") + + config = BrowserTakeoverConfig.from_environment() + + assert config is not None + assert config.public_url == "https://blacki.example.ts.net/browser-takeover" + assert config.public_origin == "https://blacki.example.ts.net" + assert config.secure_cookie is True + assert config.ttl_seconds == 420 + assert config.stream_port == 9323 + + +def test_loopback_http_is_allowed() -> None: + config = BrowserTakeoverConfig("http://127.0.0.1:8080/browser-takeover") + + config.validate() + + assert config.secure_cookie is False + + +@pytest.mark.parametrize( + ("url", "message"), + [ + ("http://example.test/browser-takeover", "must use HTTPS"), + ("https:///browser-takeover", "absolute URL"), + ( + "https://user:pass@example.test/browser-takeover", + "without credentials", + ), + ("https://example.test/browser-takeover?x=1", "query or fragment"), + ("https://example.test/takeover", "must end with /browser-takeover"), + ], +) +def test_unsafe_urls_are_rejected(url: str, message: str) -> None: + with pytest.raises(BrowserTakeoverConfigurationError, match=message): + BrowserTakeoverConfig(url).validate() + + +@pytest.mark.parametrize( + "config", + [ + BrowserTakeoverConfig("https://example.test/browser-takeover", ttl_seconds=59), + BrowserTakeoverConfig("https://example.test/browser-takeover", ttl_seconds=901), + BrowserTakeoverConfig("https://example.test/browser-takeover", stream_port=80), + BrowserTakeoverConfig( + "https://example.test/browser-takeover", stream_port=70000 + ), + ], +) +def test_unsafe_limits_are_rejected(config: BrowserTakeoverConfig) -> None: + with pytest.raises(BrowserTakeoverConfigurationError): + config.validate() + + +def test_non_integer_environment_values_are_rejected( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "BROWSER_TAKEOVER_PUBLIC_URL", + "https://example.test/browser-takeover", + ) + monkeypatch.setenv("BROWSER_TAKEOVER_TTL_SECONDS", "five") + + with pytest.raises(BrowserTakeoverConfigurationError, match="must be integers"): + BrowserTakeoverConfig.from_environment() diff --git a/tests/browser_takeover/test_routes.py b/tests/browser_takeover/test_routes.py new file mode 100644 index 0000000..24b8355 --- /dev/null +++ b/tests/browser_takeover/test_routes.py @@ -0,0 +1,611 @@ +"""Tests for the private takeover page and WebSocket proxy.""" + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.websockets import WebSocketDisconnect +from websockets.exceptions import ConnectionClosed + +from blacki.browser_takeover.config import BrowserTakeoverConfig +from blacki.browser_takeover.routes import ( + COOKIE_NAME, + _safe_server_message, + _valid_client_message, + create_browser_takeover_router, +) + + +def _app() -> FastAPI: + app = FastAPI() + app.include_router(create_browser_takeover_router()) + return app + + +def _service() -> SimpleNamespace: + return SimpleNamespace( + config=BrowserTakeoverConfig( + "http://127.0.0.1/browser-takeover", + ttl_seconds=60, + ), + redeem=AsyncMock(return_value="browser-cookie"), + complete=AsyncMock(return_value=True), + authorize=AsyncMock(return_value=None), + ) + + +def _authorized_session() -> SimpleNamespace: + return SimpleNamespace( + upstream_url="wss://sandbox.internal/private", + upstream_headers={"X-Route": "secret"}, + expected_origin="https://accounts.example.test", + ) + + +def test_page_is_hidden_when_takeover_is_disabled() -> None: + with patch( + "blacki.browser_takeover.routes.get_browser_takeover_service", + return_value=None, + ): + response = TestClient(_app()).get("/browser-takeover") + + assert response.status_code == 404 + + +def test_page_has_private_security_headers_and_human_controls() -> None: + service = _service() + with patch("blacki.browser_takeover.routes._service", return_value=service): + response = TestClient(_app()).get("/browser-takeover") + + assert response.status_code == 200 + assert "Private browser control" in response.text + assert "Expected site: verifying..." in response.text + assert "verify this site before typing sensitive data" in response.text + assert "input_touch" in response.text + assert "mouseWheel" in response.text + assert response.headers["cache-control"] == "no-store, max-age=0" + assert response.headers["referrer-policy"] == "no-referrer" + assert response.headers["x-frame-options"] == "DENY" + + +def test_redeem_is_single_origin_and_sets_http_only_cookie() -> None: + service = _service() + client = TestClient(_app()) + with patch("blacki.browser_takeover.routes._service", return_value=service): + rejected = client.post( + "/browser-takeover/redeem", + json={"token": "one-time"}, + headers={"origin": "https://attacker.test"}, + ) + response = client.post( + "/browser-takeover/redeem", + json={"token": "one-time"}, + headers={"origin": "http://127.0.0.1"}, + ) + + assert rejected.status_code == 404 + assert response.status_code == 200 + assert response.json() == {"status": "ready"} + assert COOKIE_NAME in response.headers["set-cookie"] + assert "HttpOnly" in response.headers["set-cookie"] + assert "SameSite=strict" in response.headers["set-cookie"] + service.redeem.assert_awaited_once_with("one-time") + + +def test_redeem_rejects_bad_json_and_expired_token() -> None: + service = _service() + service.redeem.return_value = None + with patch("blacki.browser_takeover.routes._service", return_value=service): + response = TestClient(_app()).post( + "/browser-takeover/redeem", + content="not-json", + headers={ + "content-type": "application/json", + "origin": "http://127.0.0.1", + }, + ) + + assert response.status_code == 404 + service.redeem.assert_awaited_once_with("") + + with patch("blacki.browser_takeover.routes._service", return_value=service): + too_large = TestClient(_app()).post( + "/browser-takeover/redeem", + content="x" * 4097, + headers={"origin": "http://127.0.0.1"}, + ) + assert too_large.status_code == 404 + + +def test_redeem_rejects_non_string_tokens() -> None: + service = _service() + service.redeem.return_value = None + with patch("blacki.browser_takeover.routes._service", return_value=service): + response = TestClient(_app()).post( + "/browser-takeover/redeem", + json={"token": 123}, + headers={"origin": "http://127.0.0.1"}, + ) + + assert response.status_code == 404 + service.redeem.assert_awaited_once_with("") + + +def test_complete_requires_origin_and_valid_cookie() -> None: + service = _service() + client = TestClient(_app()) + client.cookies.set(COOKIE_NAME, "browser-cookie") + with patch("blacki.browser_takeover.routes._service", return_value=service): + rejected = client.post( + "/browser-takeover/complete", + headers={"origin": "https://attacker.test"}, + ) + response = client.post( + "/browser-takeover/complete", + headers={"origin": "http://127.0.0.1"}, + ) + + assert rejected.status_code == 404 + assert response.status_code == 200 + assert response.json() == {"status": "complete"} + assert "Max-Age=0" in response.headers["set-cookie"] + service.complete.assert_awaited_once_with("browser-cookie") + + +def test_client_protocol_is_allowlisted_and_shape_validated() -> None: + assert ( + _valid_client_message( + '{"type":"input_keyboard","eventType":"keyDown","key":"x","code":"KeyX"}' + ) + is True + ) + assert ( + _valid_client_message( + '{"type":"input_mouse","eventType":"mouseWheel","x":10,"y":20,"deltaY":100}' + ) + is True + ) + assert ( + _valid_client_message( + '{"type":"input_mouse","eventType":"mousePressed","x":10,"y":20}' + ) + is True + ) + assert ( + _valid_client_message( + '{"type":"input_touch","eventType":"touchStart",' + '"touchPoints":[{"x":10,"y":20,"id":1}]}' + ) + is True + ) + assert _valid_client_message('{"type":"ack","seq":3}') is True + assert _valid_client_message('{"type":"config","maxFps":12,"pacing":"ack"}') is True + + assert _valid_client_message('{"type":"status","connected":true}') is False + assert _valid_client_message('{"type":"run_command"}') is False + assert ( + _valid_client_message( + '{"type":"input_keyboard","eventType":"keyDown",' + '"key":"x","code":"KeyX","command":"rm"}' + ) + is False + ) + assert _valid_client_message('{"type":"config","maxFps":121}') is False + assert _valid_client_message("not-json") is False + assert _valid_client_message("x" * (16 * 1024 + 1)) is False + + +@pytest.mark.parametrize( + "payload", + [ + {"type": "ack", "seq": 3, "extra": True}, + {"type": "config"}, + {"type": "config", "extra": True}, + { + "type": "input_keyboard", + "eventType": "unsupported", + "key": "x", + }, + { + "type": "input_keyboard", + "eventType": "keyDown", + "key": "x", + "modifiers": 16, + }, + {"type": "input_keyboard", "eventType": "char", "text": ""}, + { + "type": "input_mouse", + "eventType": "unsupported", + "x": 1, + "y": 1, + }, + { + "type": "input_mouse", + "eventType": "mousePressed", + "x": 1, + "y": 1, + "extra": True, + }, + { + "type": "input_mouse", + "eventType": "mousePressed", + "x": "bad", + "y": 1, + }, + { + "type": "input_mouse", + "eventType": "mousePressed", + "x": 1, + "y": 1, + "modifiers": 16, + }, + { + "type": "input_mouse", + "eventType": "mousePressed", + "x": 1, + "y": 1, + "button": "bad", + }, + { + "type": "input_mouse", + "eventType": "mousePressed", + "x": 1, + "y": 1, + "clickCount": 4, + }, + { + "type": "input_mouse", + "eventType": "mouseWheel", + "x": 1, + "y": 1, + }, + { + "type": "input_touch", + "eventType": "touchStart", + "touchPoints": [1], + }, + { + "type": "input_touch", + "eventType": "touchStart", + "touchPoints": [{"x": 1, "y": 1, "id": 1, "extra": True}], + }, + { + "type": "input_touch", + "eventType": "touchStart", + "touchPoints": [{"x": "bad", "y": 1}], + }, + { + "type": "input_touch", + "eventType": "unsupported", + "touchPoints": [], + }, + { + "type": "input_touch", + "eventType": "touchStart", + "touchPoints": [], + }, + { + "type": "input_touch", + "eventType": "touchStart", + "touchPoints": "bad", + }, + { + "type": "input_touch", + "eventType": "touchStart", + "touchPoints": [{"x": 0, "y": 0}] * 11, + }, + { + "type": "input_touch", + "eventType": "touchMove", + "touchPoints": [], + }, + { + "type": "input_touch", + "eventType": "touchStart", + "touchPoints": [], + "extra": True, + }, + [], + ], +) +def test_client_protocol_rejects_malformed_shapes(payload: object) -> None: + assert _valid_client_message(json.dumps(payload)) is False + + +def test_server_protocol_drops_sensitive_noise_and_sanitizes_navigation() -> None: + safe_url = _safe_server_message( + '{"type":"url","url":"https://EXAMPLE.test:443/login?token=secret#step"}' + ) + assert safe_url is not None + assert json.loads(safe_url) == { + "type": "navigation", + "origin": "https://example.test", + } + assert "secret" not in safe_url + + safe_frame = _safe_server_message( + '{"type":"frame","seq":7,"data":"abc",' + '"metadata":{"deviceWidth":1280,"deviceHeight":720,"evil":"secret"}}' + ) + assert safe_frame is not None + assert json.loads(safe_frame) == { + "type": "frame", + "seq": 7, + "data": "abc", + "metadata": {"deviceWidth": 1280, "deviceHeight": 720}, + } + + assert _safe_server_message('{"type":"console","text":"password"}') is None + assert _safe_server_message('{"type":"tabs","tabs":[]}') is None + assert _safe_server_message("[]") is None + assert _safe_server_message('{"type":"frame","data":123}') is None + frame_without_optional_data = _safe_server_message( + '{"type":"frame","seq":-1,"data":"abc"}' + ) + assert frame_without_optional_data is not None + assert json.loads(frame_without_optional_data) == { + "type": "frame", + "data": "abc", + } + frame_without_safe_metadata = _safe_server_message( + '{"type":"frame","data":"abc","metadata":{"evil":"secret"}}' + ) + assert frame_without_safe_metadata is not None + assert json.loads(frame_without_safe_metadata) == { + "type": "frame", + "data": "abc", + } + assert _safe_server_message('{"type":"status","connected":"yes"}') is None + assert _safe_server_message('{"type":"url","url":123}') is None + assert ( + _safe_server_message('{"type":"url","url":"https://example.test:bad"}') is None + ) + assert _safe_server_message('{"type":"url","url":"ftp://example.test"}') is None + ipv6_navigation = _safe_server_message( + '{"type":"url","url":"https://[::1]:443/login"}' + ) + assert ipv6_navigation is not None + assert json.loads(ipv6_navigation) == { + "type": "navigation", + "origin": "https://[::1]", + } + assert _safe_server_message("not-json") is None + + +def test_websocket_rejects_missing_cookie() -> None: + service = _service() + with ( + patch("blacki.browser_takeover.routes._service", return_value=service), + pytest.raises(WebSocketDisconnect) as caught, + TestClient(_app()).websocket_connect("/browser-takeover/ws"), + ): + pass + + assert caught.value.code == 4401 + + +def test_websocket_rejects_missing_cookie_with_matching_origin() -> None: + service = _service() + with ( + patch("blacki.browser_takeover.routes._service", return_value=service), + pytest.raises(WebSocketDisconnect) as caught, + TestClient(_app()).websocket_connect( + "/browser-takeover/ws", + headers={"origin": "http://127.0.0.1"}, + ), + ): + pass + + assert caught.value.code == 4401 + service.authorize.assert_awaited_once_with(None) + + +def test_websocket_is_hidden_when_takeover_is_disabled() -> None: + with ( + patch("blacki.browser_takeover.routes._service", return_value=None), + pytest.raises(WebSocketDisconnect) as caught, + TestClient(_app()).websocket_connect("/browser-takeover/ws"), + ): + pass + + assert caught.value.code == 4401 + + +class _Upstream: + def __init__(self) -> None: + self.sent: list[str] = [] + + async def __aenter__(self) -> "_Upstream": + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + def __aiter__(self) -> "_Upstream": + return self + + async def __anext__(self) -> str: + if self.sent: + raise StopAsyncIteration + self.sent.append("upstream-started") + return '{"type":"status","connected":true}' + + async def send(self, message: str) -> None: + self.sent.append(message) + + +class _InteractiveUpstream: + def __init__(self) -> None: + self.sent: list[str] = [] + self._message_received = asyncio.Event() + self._yielded_frame = False + + async def __aenter__(self) -> "_InteractiveUpstream": + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + def __aiter__(self) -> "_InteractiveUpstream": + return self + + async def __anext__(self) -> bytes: + await self._message_received.wait() + if self._yielded_frame: + raise StopAsyncIteration + self._yielded_frame = True + return b"frame" + + async def send(self, message: str) -> None: + self.sent.append(message) + self._message_received.set() + + +class _TwoMessageUpstream: + def __init__(self) -> None: + self._messages = [ + '{"type":"console","text":"secret"}', + '{"type":"status","connected":true}', + ] + + async def __aenter__(self) -> "_TwoMessageUpstream": + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + def __aiter__(self) -> "_TwoMessageUpstream": + return self + + async def __anext__(self) -> str: + if not self._messages: + raise StopAsyncIteration + return self._messages.pop(0) + + async def send(self, _message: str) -> None: + return None + + +def test_websocket_proxies_authorized_stream_without_exposing_endpoint() -> None: + service = _service() + service.authorize.return_value = _authorized_session() + upstream = _Upstream() + client = TestClient(_app()) + client.cookies.set(COOKIE_NAME, "browser-cookie") + with ( + patch("blacki.browser_takeover.routes._service", return_value=service), + patch("blacki.browser_takeover.routes.connect", return_value=upstream) as dial, + client.websocket_connect( + "/browser-takeover/ws", + headers={"origin": "http://127.0.0.1"}, + ) as websocket, + ): + assert websocket.receive_json() == { + "type": "takeover_meta", + "expectedOrigin": "https://accounts.example.test", + } + assert websocket.receive_text() == '{"type":"status","connected":true}' + + dial.assert_called_once_with( + "wss://sandbox.internal/private", + additional_headers={"X-Route": "secret"}, + max_size=8 * 1024 * 1024, + ) + service.authorize.assert_awaited_once_with("browser-cookie") + service.complete.assert_not_awaited() + + +def test_websocket_drops_unsafe_server_messages_before_safe_messages() -> None: + service = _service() + service.authorize.return_value = _authorized_session() + upstream = _TwoMessageUpstream() + client = TestClient(_app()) + client.cookies.set(COOKIE_NAME, "browser-cookie") + with ( + patch("blacki.browser_takeover.routes._service", return_value=service), + patch("blacki.browser_takeover.routes.connect", return_value=upstream), + client.websocket_connect( + "/browser-takeover/ws", + headers={"origin": "http://127.0.0.1"}, + ) as websocket, + ): + assert websocket.receive_json()["type"] == "takeover_meta" + assert websocket.receive_json() == { + "type": "status", + "connected": True, + } + + service.complete.assert_not_awaited() + + +def test_websocket_forwards_valid_input_and_binary_frames() -> None: + service = _service() + service.authorize.return_value = _authorized_session() + upstream = _InteractiveUpstream() + client = TestClient(_app()) + client.cookies.set(COOKIE_NAME, "browser-cookie") + message = '{"type":"input_keyboard","eventType":"keyDown","key":"x","code":"KeyX"}' + with ( + patch("blacki.browser_takeover.routes._service", return_value=service), + patch("blacki.browser_takeover.routes.connect", return_value=upstream), + client.websocket_connect( + "/browser-takeover/ws", + headers={"origin": "http://127.0.0.1"}, + ) as websocket, + ): + assert websocket.receive_json()["type"] == "takeover_meta" + websocket.send_text('{"type":"run_command"}') + websocket.send_text(message) + assert websocket.receive_bytes() == b"frame" + + assert upstream.sent == [message] + service.complete.assert_not_awaited() + + +def test_websocket_ignores_upstream_disconnect() -> None: + service = _service() + service.authorize.return_value = _authorized_session() + client = TestClient(_app()) + client.cookies.set(COOKIE_NAME, "browser-cookie") + with ( + patch("blacki.browser_takeover.routes._service", return_value=service), + patch( + "blacki.browser_takeover.routes.connect", + side_effect=ConnectionClosed(None, None), + ), + client.websocket_connect( + "/browser-takeover/ws", + headers={"origin": "http://127.0.0.1"}, + ), + ): + pass + + service.complete.assert_not_awaited() + + +def test_websocket_closes_on_unexpected_upstream_error() -> None: + service = _service() + service.authorize.return_value = _authorized_session() + client = TestClient(_app()) + client.cookies.set(COOKIE_NAME, "browser-cookie") + with ( + patch("blacki.browser_takeover.routes._service", return_value=service), + patch( + "blacki.browser_takeover.routes.connect", + side_effect=RuntimeError("upstream unavailable"), + ), + client.websocket_connect( + "/browser-takeover/ws", + headers={"origin": "http://127.0.0.1"}, + ) as websocket, + ): + assert websocket.receive_json()["type"] == "takeover_meta" + with pytest.raises(WebSocketDisconnect) as caught: + websocket.receive_text() + + assert caught.value.code == 1011 + service.complete.assert_not_awaited() diff --git a/tests/browser_takeover/test_service.py b/tests/browser_takeover/test_service.py new file mode 100644 index 0000000..7018ae7 --- /dev/null +++ b/tests/browser_takeover/test_service.py @@ -0,0 +1,520 @@ +"""Tests for one-time browser takeover leases.""" + +import asyncio +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import AsyncMock, patch + +import pytest +from opensandbox.models.execd import RunCommandOpts + +from blacki.browser_takeover.config import BrowserTakeoverConfig +from blacki.browser_takeover.service import ( + BrowserTakeoverError, + BrowserTakeoverLease, + BrowserTakeoverService, + _origin, + _websocket_url, + get_browser_takeover_service, + reset_browser_takeover_service, +) + + +class FakeCommands: + def __init__( + self, + error: object | None = None, + raise_on: str | None = None, + ) -> None: + self.error = error + self.raise_on = raise_on + self.calls: list[tuple[str, RunCommandOpts]] = [] + + async def run(self, command: str, *, opts: RunCommandOpts) -> SimpleNamespace: + self.calls.append((command, opts)) + if command == self.raise_on: + raise RuntimeError("sandbox stop failed") + return SimpleNamespace(error=self.error) + + +class FakeSandbox: + def __init__( + self, + error: object | None = None, + raise_on: str | None = None, + ) -> None: + self.commands = FakeCommands(error, raise_on=raise_on) + + async def get_endpoint(self, port: int) -> SimpleNamespace: + assert port == 9223 + return SimpleNamespace( + endpoint="https://sandbox-proxy.test/route", + headers={"X-Sandbox-Route": "private"}, + ) + + +class EndpointFailureSandbox(FakeSandbox): + async def get_endpoint(self, port: int) -> SimpleNamespace: + del port + raise RuntimeError("endpoint lookup failed") + + +@pytest.fixture +def config() -> BrowserTakeoverConfig: + return BrowserTakeoverConfig( + "https://blacki.example.ts.net/browser-takeover", + ttl_seconds=60, + ) + + +@pytest.fixture +def state() -> dict[str, str]: + return { + "telegram_chat_type": "private", + "telegram_chat_id": "42", + "temp:telegram_sender_user_id": "42", + } + + +async def _create( + service: BrowserTakeoverService, + state: dict[str, str], + sandbox: FakeSandbox | None = None, +) -> tuple[BrowserTakeoverLease, FakeSandbox]: + sandbox = sandbox or FakeSandbox() + manager = SimpleNamespace( + get_or_create_sandbox=AsyncMock( + return_value={"sandbox": sandbox, "error": None} + ) + ) + with patch( + "blacki.browser_takeover.service.get_sandbox_manager", + return_value=manager, + ): + lease = await service.create( + login_url="https://accounts.example.test/login", + state=state, + ) + return lease, sandbox + + +async def test_create_redeem_complete_wait_and_close( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + lease, sandbox = await _create(service, state) + link_token = lease.takeover_url.rsplit("#", 1)[1] + + command, opts = sandbox.commands.calls[0] + assert "accounts.example.test" not in command + assert opts.envs is not None + assert opts.envs["BLACKI_BROWSER_TAKEOVER_URL"].endswith("/login") + assert lease.takeover_url.startswith(config.public_url + "#") + + browser_token = await service.redeem(link_token) + assert browser_token is not None + assert await service.redeem("") is None + assert await service.redeem(link_token) is None + assert await service.authorize(None) is None + session = await service.authorize(browser_token) + assert session is not None + assert session.upstream_url == "wss://sandbox-proxy.test/route" + assert session.upstream_headers == {"X-Sandbox-Route": "private"} + assert session.expected_origin == "https://accounts.example.test" + assert await service.complete(browser_token) is True + assert await service.wait(lease) is True + + await service.close(lease) + await service.close(lease) + + assert sandbox.commands.calls[-1][0] == "agent-browser stream disable" + assert await service.authorize(browser_token) is None + assert await service.complete(browser_token) is False + assert await service.wait(lease) is False + + +async def test_wait_expires_and_close_all_wakes_without_success( + config: BrowserTakeoverConfig, + state: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = BrowserTakeoverService(config) + lease, expired_sandbox = await _create(service, state) + with monkeypatch.context() as context: + context.setattr( + "blacki.browser_takeover.service.time.monotonic", + lambda: 10e9, + ) + assert await service.wait(lease) is False + lease2, _ = await _create(service, state) + assert await service.redeem(lease.takeover_url.rsplit("#", 1)[1]) is None + + assert any( + command == "agent-browser stream disable" + for command, _ in expired_sandbox.commands.calls + ) + + waiter = asyncio.create_task(service.wait(lease2)) + await asyncio.sleep(0) + await service.close_all() + + assert await waiter is False + + +async def test_expired_redeem_wakes_waiter_without_success( + config: BrowserTakeoverConfig, + state: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = [100.0] + monkeypatch.setattr( + "blacki.browser_takeover.service.time.monotonic", + lambda: now[0], + ) + service = BrowserTakeoverService(config) + lease, _ = await _create(service, state) + waiter = asyncio.create_task(service.wait(lease)) + await asyncio.sleep(0) + + now[0] = 10e9 + assert await service.redeem(lease.takeover_url.rsplit("#", 1)[1]) is None + assert await waiter is False + + +async def test_expired_authorize_wakes_waiter_without_success( + config: BrowserTakeoverConfig, + state: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = [100.0] + monkeypatch.setattr( + "blacki.browser_takeover.service.time.monotonic", + lambda: now[0], + ) + service = BrowserTakeoverService(config) + lease, _ = await _create(service, state) + browser_token = await service.redeem(lease.takeover_url.rsplit("#", 1)[1]) + assert browser_token is not None + waiter = asyncio.create_task(service.wait(lease)) + await asyncio.sleep(0) + + now[0] = 10e9 + assert await service.authorize(browser_token) is None + assert await waiter is False + + +async def test_replacing_expired_takeover_wakes_old_waiter_without_success( + config: BrowserTakeoverConfig, + state: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + now = [100.0] + monkeypatch.setattr( + "blacki.browser_takeover.service.time.monotonic", + lambda: now[0], + ) + service = BrowserTakeoverService(config) + lease, _ = await _create(service, state) + waiter = asyncio.create_task(service.wait(lease)) + await asyncio.sleep(0) + + now[0] = 10e9 + replacement, _ = await _create(service, state) + + assert await waiter is False + await service.close(replacement) + + +@pytest.mark.parametrize( + "bad_state", + [ + {}, + { + "telegram_chat_type": "group", + "telegram_chat_id": "42", + "temp:telegram_sender_user_id": "42", + }, + { + "telegram_chat_type": "private", + "telegram_chat_id": "42", + "temp:telegram_sender_user_id": "7", + }, + ], +) +async def test_create_requires_matching_private_telegram_identity( + config: BrowserTakeoverConfig, + bad_state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + + with pytest.raises(BrowserTakeoverError): + await service.create(login_url="https://example.test/login", state=bad_state) + + +@pytest.mark.parametrize( + "url", + [ + "http://example.test/login", + "not-a-url", + "https://user:pass@example.test", + "https://example.test:bad/login", + ], +) +async def test_create_rejects_unsafe_login_url( + config: BrowserTakeoverConfig, + state: dict[str, str], + url: str, +) -> None: + service = BrowserTakeoverService(config) + + with pytest.raises(BrowserTakeoverError): + await service.create(login_url=url, state=state) + + +async def test_create_handles_missing_or_failed_sandbox( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + unavailable = SimpleNamespace( + get_or_create_sandbox=AsyncMock( + return_value={"sandbox": None, "error": "unavailable"} + ) + ) + with ( + patch( + "blacki.browser_takeover.service.get_sandbox_manager", + return_value=unavailable, + ), + pytest.raises(BrowserTakeoverError, match="unavailable"), + ): + await service.create(login_url="https://example.test/login", state=state) + + with pytest.raises(BrowserTakeoverError, match="could not start"): + await _create(service, state, FakeSandbox(error=object())) + + +async def test_create_wraps_unexpected_start_errors( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + manager = SimpleNamespace( + get_or_create_sandbox=AsyncMock(side_effect=RuntimeError("sandbox failed")) + ) + + with ( + patch( + "blacki.browser_takeover.service.get_sandbox_manager", + return_value=manager, + ), + pytest.raises(BrowserTakeoverError, match="could not start"), + ): + await service.create(login_url="https://example.test/login", state=state) + + lease, _ = await _create(service, state) + await service.close(lease) + + +async def test_create_stops_stream_when_endpoint_lookup_fails( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + sandbox = EndpointFailureSandbox() + manager = SimpleNamespace( + get_or_create_sandbox=AsyncMock( + return_value={"sandbox": sandbox, "error": None} + ) + ) + + with ( + patch( + "blacki.browser_takeover.service.get_sandbox_manager", + return_value=manager, + ), + pytest.raises(BrowserTakeoverError, match="could not start"), + ): + await service.create(login_url="https://example.test/login", state=state) + + assert [command for command, _ in sandbox.commands.calls] == [ + 'agent-browser open "$BLACKI_BROWSER_TAKEOVER_URL" && ' + 'agent-browser stream enable --port "$BLACKI_BROWSER_STREAM_PORT"', + "agent-browser stream disable", + ] + + +async def test_cancelled_start_releases_owner_reservation( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + cancelled = SimpleNamespace( + get_or_create_sandbox=AsyncMock(side_effect=asyncio.CancelledError) + ) + with ( + patch( + "blacki.browser_takeover.service.get_sandbox_manager", + return_value=cancelled, + ), + pytest.raises(asyncio.CancelledError), + ): + await service.create(login_url="https://example.test/login", state=state) + + lease, _ = await _create(service, state) + await service.close(lease) + + +async def test_only_one_takeover_can_run_per_conversation( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + lease, _ = await _create(service, state) + + with pytest.raises(BrowserTakeoverError, match="already active"): + await _create(service, state) + + await service.close(lease) + + +async def test_only_one_takeover_can_start_per_conversation( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + started = asyncio.Event() + release = asyncio.Event() + sandbox = FakeSandbox() + + async def get_or_create_sandbox(_state: Any) -> dict[str, Any]: + started.set() + await release.wait() + return {"sandbox": sandbox, "error": None} + + manager = SimpleNamespace(get_or_create_sandbox=get_or_create_sandbox) + with patch( + "blacki.browser_takeover.service.get_sandbox_manager", + return_value=manager, + ): + first = asyncio.create_task( + service.create( + login_url="https://example.test/login", + state=state, + ) + ) + await started.wait() + with pytest.raises(BrowserTakeoverError, match="already starting"): + await service.create(login_url="https://example.test/login", state=state) + release.set() + lease = await first + + await service.close(lease) + + +async def test_authorize_removes_expired_sessions( + config: BrowserTakeoverConfig, + state: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = BrowserTakeoverService(config) + lease, sandbox = await _create(service, state) + browser_token = await service.redeem(lease.takeover_url.rsplit("#", 1)[1]) + assert browser_token is not None + monkeypatch.setattr("blacki.browser_takeover.service.time.monotonic", lambda: 10e9) + + assert await service.authorize(browser_token) is None + assert [command for command, _ in sandbox.commands.calls].count( + "agent-browser stream disable" + ) == 1 + + +async def test_redeem_removes_expired_sessions( + config: BrowserTakeoverConfig, + state: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + service = BrowserTakeoverService(config) + lease, sandbox = await _create(service, state) + monkeypatch.setattr("blacki.browser_takeover.service.time.monotonic", lambda: 10e9) + + assert await service.redeem(lease.takeover_url.rsplit("#", 1)[1]) is None + assert [command for command, _ in sandbox.commands.calls].count( + "agent-browser stream disable" + ) == 1 + + +async def test_redeem_rejects_non_string_tokens( + config: BrowserTakeoverConfig, +) -> None: + service = BrowserTakeoverService(config) + + assert await service.redeem(cast(str, 123)) is None + assert await service.redeem("x" * 257) is None + + +async def test_close_swallows_stream_stop_errors( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + sandbox = FakeSandbox(raise_on="agent-browser stream disable") + lease, _ = await _create(service, state, sandbox) + + await service.close(lease) + + +async def test_close_all_deduplicates_sandboxes_and_swallows_errors( + config: BrowserTakeoverConfig, + state: dict[str, str], +) -> None: + service = BrowserTakeoverService(config) + sandbox = FakeSandbox(raise_on="agent-browser stream disable") + other_state = { + **state, + "telegram_chat_id": "43", + "temp:telegram_sender_user_id": "43", + } + await _create(service, state, sandbox) + await _create(service, other_state, sandbox) + + await service.close_all() + + stop_calls = [ + command + for command, _ in sandbox.commands.calls + if command == "agent-browser stream disable" + ] + assert len(stop_calls) == 1 + + +def test_websocket_endpoint_normalization() -> None: + assert _websocket_url("sandbox:9223") == "ws://sandbox:9223" + assert _websocket_url("http://sandbox/route") == "ws://sandbox/route" + assert _websocket_url("wss://sandbox/route") == "wss://sandbox/route" + + +def test_login_origin_normalizes_ipv6_hosts() -> None: + assert _origin("https://[::1]/login") == "https://[::1]" + + +def test_login_origin_rejects_missing_hosts() -> None: + with pytest.raises(BrowserTakeoverError, match="valid login host"): + _origin("https://:443/login") + + +async def test_process_service_lifecycle(monkeypatch: pytest.MonkeyPatch) -> None: + await reset_browser_takeover_service() + monkeypatch.delenv("BROWSER_TAKEOVER_PUBLIC_URL", raising=False) + assert get_browser_takeover_service() is None + + monkeypatch.setenv( + "BROWSER_TAKEOVER_PUBLIC_URL", + "https://blacki.example.ts.net/browser-takeover", + ) + first = get_browser_takeover_service() + assert first is not None + assert get_browser_takeover_service() is first + + await reset_browser_takeover_service() diff --git a/tests/browser_takeover/test_tools.py b/tests/browser_takeover/test_tools.py new file mode 100644 index 0000000..d904c1b --- /dev/null +++ b/tests/browser_takeover/test_tools.py @@ -0,0 +1,156 @@ +"""Tests for the model-blind Telegram takeover tool.""" + +from types import SimpleNamespace +from typing import cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from google.adk.tools import ToolContext + +from blacki.browser_takeover import BrowserTakeoverError, BrowserTakeoverLease +from blacki.tools.browser_takeover import _send_takeover_link, start_browser_takeover + + +def _context() -> ToolContext: + return cast( + ToolContext, + SimpleNamespace( + state={ + "telegram_chat_id": "42", + "telegram_thread_id": "9", + "telegram_chat_type": "private", + "temp:telegram_sender_user_id": "42", + }, + ), + ) + + +async def test_tool_returns_status_without_leaking_link() -> None: + lease = BrowserTakeoverLease( + session_id="session", + takeover_url="https://private.test/browser#secret-link-token", + expires_in_seconds=300, + ) + service = SimpleNamespace( + create=AsyncMock(return_value=lease), + wait=AsyncMock(return_value=True), + close=AsyncMock(), + ) + with ( + patch( + "blacki.tools.browser_takeover.get_browser_takeover_service", + return_value=service, + ), + patch( + "blacki.tools.browser_takeover._send_takeover_link", + new=AsyncMock(), + ) as send, + ): + result = await start_browser_takeover( + "https://accounts.example.test/login", + _context(), + ) + + assert result["status"] == "success" + assert "secret-link-token" not in repr(result) + send.assert_awaited_once_with(lease.takeover_url, 300, _context().state) + service.close.assert_awaited_once_with(lease) + + +async def test_tool_handles_unavailable_create_expiry_and_delivery_failure() -> None: + with patch( + "blacki.tools.browser_takeover.get_browser_takeover_service", + return_value=None, + ): + assert (await start_browser_takeover("https://example.test", _context()))[ + "status" + ] == "unavailable" + + create_failure = SimpleNamespace( + create=AsyncMock(side_effect=BrowserTakeoverError("safe failure")) + ) + with patch( + "blacki.tools.browser_takeover.get_browser_takeover_service", + return_value=create_failure, + ): + result = await start_browser_takeover("https://example.test", _context()) + assert result == {"status": "error", "message": "safe failure"} + + lease = BrowserTakeoverLease("session", "https://private.test/#token", 300) + expired = SimpleNamespace( + create=AsyncMock(return_value=lease), + wait=AsyncMock(return_value=False), + close=AsyncMock(), + ) + with ( + patch( + "blacki.tools.browser_takeover.get_browser_takeover_service", + return_value=expired, + ), + patch( + "blacki.tools.browser_takeover._send_takeover_link", + new=AsyncMock(), + ), + ): + result = await start_browser_takeover("https://example.test", _context()) + assert result["status"] == "expired" + expired.close.assert_awaited_once_with(lease) + + delivery_failure = SimpleNamespace( + create=AsyncMock(return_value=lease), + wait=AsyncMock(), + close=AsyncMock(), + ) + with ( + patch( + "blacki.tools.browser_takeover.get_browser_takeover_service", + return_value=delivery_failure, + ), + patch( + "blacki.tools.browser_takeover._send_takeover_link", + new=AsyncMock(side_effect=RuntimeError("secret provider failure")), + ), + ): + result = await start_browser_takeover("https://example.test", _context()) + assert result["status"] == "error" + assert "secret provider failure" not in repr(result) + delivery_failure.close.assert_awaited_once_with(lease) + + +async def test_link_is_sent_directly_as_protected_telegram_content( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TELEGRAM_BOT_TOKEN", "bot-token") + api = AsyncMock() + context_manager = MagicMock() + context_manager.__aenter__ = AsyncMock(return_value=api) + context_manager.__aexit__ = AsyncMock(return_value=None) + with patch( + "blacki.tools.browser_takeover.TelegramApiClient", + return_value=context_manager, + ): + await _send_takeover_link( + "https://private.test/browser#token", + 300, + _context().state, + ) + + kwargs = api.send_message.await_args.kwargs + assert kwargs["chat_id"] == 42 + assert kwargs["message_thread_id"] == 9 + assert kwargs["protect_content"] is True + assert "#token" not in kwargs["text"] + assert kwargs["reply_markup"].inline_keyboard[0][0].url.endswith("#token") + + +async def test_link_delivery_requires_telegram_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("TELEGRAM_BOT_TOKEN", raising=False) + + try: + await _send_takeover_link("https://private.test/#token", 60, _context().state) + except BrowserTakeoverError as exc: + assert str(exc) == "Telegram is not configured" + else: + raise AssertionError("Expected missing Telegram configuration to fail") diff --git a/tests/test_privacy.py b/tests/test_privacy.py index 1272cdc..18cc967 100644 --- a/tests/test_privacy.py +++ b/tests/test_privacy.py @@ -61,6 +61,7 @@ def test_private_tool_identification_uses_zepto_and_gmail_prefix() -> None: assert is_private_tool(_tool("send_text_to_speech")) is True assert is_private_tool(_tool("search_products")) is False assert is_private_tool(_tool("get_health_summary")) is True + assert is_private_tool(_tool("start_browser_takeover")) is True def test_configure_zepto_privacy_is_explicit_and_forces_safe_values( @@ -137,6 +138,17 @@ def test_kokoro_tts_enables_content_redaction( assert os.environ["OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"] == "false" +def test_browser_takeover_enables_content_redaction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "BROWSER_TAKEOVER_PUBLIC_URL", + "https://blacki.example.ts.net/browser-takeover", + ) + + assert private_tool_privacy_enabled() is True + + def test_secure_zepto_app_removes_content_logging_plugin( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/test_registry.py b/tests/test_registry.py index 70024b0..7fb9298 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -40,6 +40,7 @@ def test_default_values(self) -> None: assert config.zepto_mcp_enabled is False assert config.zepto_mcp_allowed_chat_ids == frozenset() assert config.gmail_config is None + assert config.browser_takeover_enabled is False def test_custom_values(self) -> None: """Should accept custom values.""" @@ -290,6 +291,28 @@ def test_invalid_r2_file_tools_degrade_cleanly( } assert "Durable R2 file tools disabled" in caplog.text + def test_browser_takeover_is_private_root_only(self) -> None: + """Browser takeover must never reach public or worker agents.""" + config = ToolConfig( + weather_enabled=False, + sandbox_enabled=True, + browser_takeover_enabled=True, + ) + + root_tools = build_tools(config, include_user_scoped_tools=True) + worker_tools = build_tools(config, include_user_scoped_tools=False) + default_tools = build_tools(config) + + assert "start_browser_takeover" in { + getattr(tool, "__name__", "") for tool in root_tools + } + assert "start_browser_takeover" not in { + getattr(tool, "__name__", "") for tool in worker_tools + } + assert "start_browser_takeover" not in { + getattr(tool, "__name__", "") for tool in default_tools + } + def test_google_health_tool_is_telegram_root_only(self) -> None: """Google Health must never reach the public or delegated worker agent.""" config = ToolConfig( @@ -445,6 +468,7 @@ def test_empty_env(self) -> None: assert config.sandbox_enabled is False assert config.skills_dir is not None assert config.zepto_mcp_enabled is False + assert config.browser_takeover_enabled is False assert config.kokoro_tts_base_url is None assert config.kokoro_tts_voice == "af_heart" assert config.zepto_mcp_config_dir == Path( @@ -511,6 +535,21 @@ def test_sandbox_enabled_from_env_false(self) -> None: config = build_tool_config_from_env() assert config.sandbox_enabled is False, f"Failed for value: {value}" + def test_browser_takeover_enabled_from_valid_public_url(self) -> None: + """A validated public URL should register the private handoff tool.""" + with patch.dict( + "os.environ", + { + "BROWSER_TAKEOVER_PUBLIC_URL": ( + "https://blacki.example.ts.net/browser-takeover" + ) + }, + clear=True, + ): + config = build_tool_config_from_env() + + assert config.browser_takeover_enabled is True + def test_skills_dir_always_set(self) -> None: """Should always set skills_dir to package skills directory.""" with patch.dict("os.environ", {}, clear=True): diff --git a/uv.lock b/uv.lock index ba8a254..0453f59 100644 --- a/uv.lock +++ b/uv.lock @@ -261,6 +261,7 @@ dependencies = [ { name = "pydantic" }, { name = "python-dotenv" }, { name = "pyyaml" }, + { name = "websockets" }, ] [package.dev-dependencies] @@ -304,6 +305,7 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.11.0,<3.0.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0.0" }, { name = "pyyaml", specifier = ">=6.0.0,<7.0.0" }, + { name = "websockets", specifier = ">=15.0.1,<16.0.0" }, ] [package.metadata.requires-dev]