diff --git a/doc/code/targets/10_3_websocket_copilot_target.ipynb b/doc/code/targets/10_3_websocket_copilot_target.ipynb index 41ba5dc18c..4965b9d6ba 100644 --- a/doc/code/targets/10_3_websocket_copilot_target.ipynb +++ b/doc/code/targets/10_3_websocket_copilot_target.ipynb @@ -3,7 +3,9 @@ { "cell_type": "markdown", "id": "0", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 0 + }, "source": [ "# WebSocket Copilot Target\n", "\n", @@ -13,13 +15,15 @@ "- `COPILOT_USERNAME` and `COPILOT_PASSWORD` environment variables\n", "- Playwright installed: `pip install playwright && playwright install chromium`\n", "\n", - "Some environments are not suited for automated authentication (e.g. they have security policies with retrieving tokens or have MFA). See the [Alternative Authentication](#alternative-authentication-with-manualcopilotauthenticator) section below." + "Some environments are not suited for automated authentication (e.g. they have security policies with retrieving tokens or have MFA). For interactive authentication compatible with MFA and Conditional Access, see [Browser Session Authentication](#browser-session-authentication). To provide a token manually, see [Alternative Authentication](#alternative-authentication-with-manualcopilotauthenticator)." ] }, { "cell_type": "markdown", "id": "1", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 0 + }, "source": [ "## Basic Usage with `PromptSendingAttack`\n", "\n", @@ -30,7 +34,9 @@ "cell_type": "code", "execution_count": null, "id": "2", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 0 + }, "outputs": [ { "name": "stdout", @@ -92,7 +98,9 @@ "cell_type": "code", "execution_count": null, "id": "4", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 0 + }, "outputs": [ { "name": "stdout", @@ -166,7 +174,6 @@ "source": [ "from pyrit.executor.attack import MultiPromptSendingAttack\n", "from pyrit.models import Message\n", - "from pyrit.output import output_attack_async\n", "from pyrit.prompt_target import WebSocketCopilotTarget\n", "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", "\n", @@ -195,6 +202,105 @@ { "cell_type": "markdown", "id": "5", + "metadata": { + "lines_to_next_cell": 0 + }, + "source": [ + "## Browser Session Authentication\n", + "\n", + "`BrowserSessionCopilotAuthenticator` captures a token from a persistent Edge session.\n", + "Complete account selection when prompted. The browser remains minimized for token\n", + "renewal and closes when the context exits.\n", + "\n", + "Captured tokens remain in memory. Supply a different `profile_path` for each persona.\n", + "\n", + "Install the optional dependency before using this authenticator:\n", + "\n", + "```bash\n", + "pip install \"pyrit[playwright]\"\n", + "```\n", + "\n", + "The authenticator uses a locally installed Microsoft Edge browser." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": { + "lines_to_next_cell": 0 + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n", + "\u001b[1m\u001b[33m ❓ ATTACK RESULT: UNDETERMINED ❓ \u001b[0m\n", + "\u001b[33m════════════════════════════════════════════════════════════════════════════════════════════════════\u001b[0m\n", + "\n", + "\u001b[1m\u001b[44m\u001b[37m Attack Summary \u001b[0m\n", + "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[1m 📋 Basic Information\u001b[0m\n", + "\u001b[36m • Objective: What is your favorite color?\u001b[0m\n", + "\u001b[36m • Attack Type: PromptSendingAttack\u001b[0m\n", + "\u001b[36m • Conversation ID: 4adc3983-5751-41d8-bc26-6cd76ea48b60\u001b[0m\n", + "\n", + "\u001b[1m ⚡ Execution Metrics\u001b[0m\n", + "\u001b[32m • Turns Executed: 1\u001b[0m\n", + "\u001b[32m • Execution Time: 17.48s\u001b[0m\n", + "\n", + "\u001b[1m 🎯 Outcome\u001b[0m\n", + "\u001b[33m • Status: ❓ UNDETERMINED\u001b[0m\n", + "\u001b[37m • Reason: No objective scorer configured\u001b[0m\n", + "\n", + "\u001b[1m\u001b[44m\u001b[37m Conversation History with Objective Target \u001b[0m\n", + "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\n", + "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[1m\u001b[34m🔹 Turn 1 - USER\u001b[0m\n", + "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[34m What is your favorite color?\u001b[0m\n", + "\n", + "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n", + "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[33m I don't have personal preferences, so I don't actually have a favorite color.\u001b[0m\n", + "\u001b[33m \u001b[0m\n", + "\u001b[33m If I had to pick one for fun, I'd choose **blue** because it's commonly associated with\u001b[0m\n", + "\u001b[33m technology, reliability, and calmness. 💙\u001b[0m\n", + "\u001b[33m \u001b[0m\n", + "\u001b[33m What's your favorite color?\u001b[0m\n", + "\n", + "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\n", + "\u001b[2m\u001b[37m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[2m\u001b[37m Report generated at: 2026-08-27 16:34:36 UTC \u001b[0m\n" + ] + } + ], + "source": [ + "from pyrit.auth import BrowserSessionCopilotAuthenticator\n", + "from pyrit.executor.attack import PromptSendingAttack\n", + "from pyrit.output import output_attack_async\n", + "from pyrit.prompt_target import WebSocketCopilotTarget\n", + "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", + "\n", + "await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True)\n", + "\n", + "objective = \"What is your favorite color?\"\n", + "\n", + "async with BrowserSessionCopilotAuthenticator() as auth:\n", + " target = WebSocketCopilotTarget(authenticator=auth)\n", + " attack = PromptSendingAttack(objective_target=target)\n", + " result = await attack.execute_async(objective=objective)\n", + " await output_attack_async(result)" + ] + }, + { + "cell_type": "markdown", + "id": "7", "metadata": {}, "source": [ "## Alternative Authentication with `ManualCopilotAuthenticator`\n", @@ -208,7 +314,7 @@ "3. Go to the Network tab.\n", "4. Filter by \"Socket\" connections or search for \"Chathub\".\n", "5. Start typing in the chat to initiate a WebSocket connection.\n", - "6. Look for the latest WebSocket connection to `substrate.office.com/m365Copilot/Chathub`.\n", + "6. Look for the latest WebSocket connection to `substrate.svc.cloud.microsoft/m365Copilot/Chathub`.\n", "7. You may find the `access_token` in the request URL or in the request payload.\n", "\n", "You can either pass the token directly or set the `COPILOT_ACCESS_TOKEN` environment variable." @@ -217,8 +323,10 @@ { "cell_type": "code", "execution_count": null, - "id": "6", - "metadata": {}, + "id": "8", + "metadata": { + "lines_to_next_cell": 0 + }, "outputs": [ { "name": "stdout", @@ -273,7 +381,7 @@ }, { "cell_type": "markdown", - "id": "7", + "id": "9", "metadata": {}, "source": [ "## Multimodal Support (Text and Images)\n", @@ -286,7 +394,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "10", "metadata": {}, "outputs": [ { @@ -360,7 +468,6 @@ "\n", "from pyrit.executor.attack import PromptSendingAttack\n", "from pyrit.models import Message, MessagePiece\n", - "from pyrit.output import output_attack_async\n", "from pyrit.prompt_target import WebSocketCopilotTarget\n", "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", "\n", @@ -400,6 +507,9 @@ } ], "metadata": { + "jupytext": { + "main_language": "python" + }, "language_info": { "codemirror_mode": { "name": "ipython", diff --git a/doc/code/targets/10_3_websocket_copilot_target.py b/doc/code/targets/10_3_websocket_copilot_target.py index b8cbb7f100..7b1789315c 100644 --- a/doc/code/targets/10_3_websocket_copilot_target.py +++ b/doc/code/targets/10_3_websocket_copilot_target.py @@ -16,7 +16,7 @@ # - `COPILOT_USERNAME` and `COPILOT_PASSWORD` environment variables # - Playwright installed: `pip install playwright && playwright install chromium` # -# Some environments are not suited for automated authentication (e.g. they have security policies with retrieving tokens or have MFA). See the [Alternative Authentication](#alternative-authentication-with-manualcopilotauthenticator) section below. +# Some environments are not suited for automated authentication (e.g. they have security policies with retrieving tokens or have MFA). For interactive authentication compatible with MFA and Conditional Access, see [Browser Session Authentication](#browser-session-authentication). To provide a token manually, see [Alternative Authentication](#alternative-authentication-with-manualcopilotauthenticator). # %% [markdown] # ## Basic Usage with `PromptSendingAttack` @@ -77,6 +77,39 @@ await output_attack_async(result) +# %% [markdown] +# ## Browser Session Authentication +# +# `BrowserSessionCopilotAuthenticator` captures a token from a persistent Edge session. +# Complete account selection when prompted. The browser remains minimized for token +# renewal and closes when the context exits. + +# Captured tokens remain in memory. Supply a different `profile_path` for each persona. +# +# Install the optional dependency before using this authenticator: +# +# ```bash +# pip install "pyrit[playwright]" +# ``` +# +# The authenticator uses a locally installed Microsoft Edge browser. +# %% +from pyrit.auth import BrowserSessionCopilotAuthenticator +from pyrit.executor.attack import PromptSendingAttack +from pyrit.output import output_attack_async +from pyrit.prompt_target import WebSocketCopilotTarget +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=True) + +objective = "What is your favorite color?" + +async with BrowserSessionCopilotAuthenticator() as auth: + target = WebSocketCopilotTarget(authenticator=auth) + attack = PromptSendingAttack(objective_target=target) + result = await attack.execute_async(objective=objective) + await output_attack_async(result) + # %% [markdown] # ## Alternative Authentication with `ManualCopilotAuthenticator` # @@ -89,7 +122,7 @@ # 3. Go to the Network tab. # 4. Filter by "Socket" connections or search for "Chathub". # 5. Start typing in the chat to initiate a WebSocket connection. -# 6. Look for the latest WebSocket connection to `substrate.office.com/m365Copilot/Chathub`. +# 6. Look for the latest WebSocket connection to `substrate.svc.cloud.microsoft/m365Copilot/Chathub`. # 7. You may find the `access_token` in the request URL or in the request payload. # # You can either pass the token directly or set the `COPILOT_ACCESS_TOKEN` environment variable. diff --git a/pyrit/auth/__init__.py b/pyrit/auth/__init__.py index e20000e0c8..7da0755b3c 100644 --- a/pyrit/auth/__init__.py +++ b/pyrit/auth/__init__.py @@ -25,6 +25,9 @@ is_azure_openai_endpoint, ) from pyrit.auth.azure_storage_auth import AzureStorageAuth + from pyrit.auth.browser_session_copilot_authenticator import ( + BrowserSessionCopilotAuthenticator, + ) from pyrit.auth.copilot_authenticator import CopilotAuthenticator from pyrit.auth.manual_copilot_authenticator import ManualCopilotAuthenticator from pyrit.auth.openai_auth import resolve_openai_auth @@ -34,6 +37,7 @@ "Authenticator": "pyrit.auth.authenticator", "AzureAuth": "pyrit.auth.azure_auth", "AzureStorageAuth": "pyrit.auth.azure_storage_auth", + "BrowserSessionCopilotAuthenticator": "pyrit.auth.browser_session_copilot_authenticator", "CopilotAuthenticator": "pyrit.auth.copilot_authenticator", "ManualCopilotAuthenticator": "pyrit.auth.manual_copilot_authenticator", "resolve_openai_auth": "pyrit.auth.openai_auth", diff --git a/pyrit/auth/browser_session_copilot_authenticator.py b/pyrit/auth/browser_session_copilot_authenticator.py new file mode 100644 index 0000000000..b87b7e43d9 --- /dev/null +++ b/pyrit/auth/browser_session_copilot_authenticator.py @@ -0,0 +1,502 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +import sys +import threading +import time +from collections.abc import Callable, Coroutine +from contextlib import AsyncExitStack +from pathlib import Path +from types import TracebackType +from typing import Any, TypeVar +from urllib.parse import parse_qs, urlparse + +import jwt +from typing_extensions import Self + +from pyrit.auth.authenticator import Authenticator +from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH + +T = TypeVar("T") + + +class BrowserSessionCopilotAuthenticator(Authenticator): + """Acquire Microsoft Copilot access token from a browser session.""" + + DEFAULT_TOKEN_CAPTURE_TIMEOUT_SECONDS = 60 + DEFAULT_EXPIRY_BUFFER_SECONDS = 300 + DEFAULT_WEBSOCKET_BASE_URL = "wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub" + DEFAULT_BROWSER_CHANNEL = "msedge" + DEFAULT_COPILOT_URL = "https://m365.cloud.microsoft/chat" + + def __init__( + self, + *, + profile_path: Path | None = None, + token_capture_timeout_seconds: int = DEFAULT_TOKEN_CAPTURE_TIMEOUT_SECONDS, + expiry_buffer_seconds: int = DEFAULT_EXPIRY_BUFFER_SECONDS, + websocket_base_url: str = DEFAULT_WEBSOCKET_BASE_URL, + browser_channel: str = DEFAULT_BROWSER_CHANNEL, + headless: bool = False, + copilot_url: str = DEFAULT_COPILOT_URL, + ) -> None: + """ + Initialize the authenticator with a persistent browser profile. + + Args: + profile_path (Path | None): Path to the persistent browser profile. If None, a default path is used. + token_capture_timeout_seconds (int): Timeout in seconds for capturing the token. Must be a positive integer. + expiry_buffer_seconds (int): Buffer time in seconds before token expiry. Must be a positive integer. + websocket_base_url (str): Base URL for the Copilot Chathub websocket connection. Must be a valid wss URL. + browser_channel (str): Browser channel to use for the session. Defaults to "msedge". + headless (bool): Whether to run the browser in headless mode. Defaults to False. + copilot_url (str): URL for the Copilot chat interface. Must be a valid HTTPS URL. + + Raises: + ValueError: If token_capture_timeout_seconds or expiry_buffer_seconds is not a positive integer. + ValueError: If websocket_base_url is not a valid wss URL with a path. + """ + self._access_token: str | None = None + self._token_fetch_lock = asyncio.Lock() + self._claims: dict[str, Any] = {} + if token_capture_timeout_seconds <= 0: + raise ValueError("token_capture_timeout_seconds must be a positive integer.") + if expiry_buffer_seconds <= 0: + raise ValueError("expiry_buffer_seconds must be a positive integer.") + + normalized_websocket_base_url = websocket_base_url.rstrip("/") + parsed_websocket_base_url = urlparse(normalized_websocket_base_url) + + if ( + parsed_websocket_base_url.scheme != "wss" + or not parsed_websocket_base_url.hostname + or not parsed_websocket_base_url.path.strip("/") + ): + raise ValueError("websocket_base_url must be a valid wss URL with a path.") + + self._profile_path = profile_path or (CONFIGURATION_DIRECTORY_PATH / "copilot_browser_profiles" / "default") + self._token_capture_timeout_seconds = token_capture_timeout_seconds + self._expiry_buffer_seconds = expiry_buffer_seconds + self._websocket_base_url = normalized_websocket_base_url + self._browser_channel = browser_channel + self._headless = headless + self._copilot_url = copilot_url + self._browser_resources: AsyncExitStack | None = None + self._browser_context: Any | None = None + self._page: Any | None = None + self._browser_loop: asyncio.AbstractEventLoop | None = None + self._browser_thread: threading.Thread | None = None + self._browser_loop_started = threading.Event() + self._browser_loop_start_error: Exception | None = None + + @property + def profile_path(self) -> Path: + """Persistent browser profile path.""" + return self._profile_path + + async def __aenter__(self) -> Self: + """ + Enter the asynchronous context manager. + + Returns: + The authenticator instance itself. + """ + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + """Exit the asynchronous context manager, closing the browser context and releasing resources.""" + await self.close_async() + + async def get_token_async(self) -> str: + """ + Return the current token, capturing a new one if necessary. + + Returns: + The current access token as a string. + + Raises: + ValueError: If captured token claims are missing, invalid, or within the expiry buffer. + """ + current_token = self._access_token + if current_token is not None and self._has_fresh_token(): + return current_token + + async with self._token_fetch_lock: + current_token = self._access_token + if current_token is not None and self._has_fresh_token(): + return current_token + + return await self._capture_and_store_token_async() + + async def refresh_token_async(self) -> str: + """ + Force a refresh of the current token, capturing a new one. + + Returns: + The newly captured access token as a string. + """ + async with self._token_fetch_lock: + return await self._capture_and_store_token_async() + + async def close_async(self) -> None: + """Close browser resources and discard authentication state.""" + async with self._token_fetch_lock: + self._access_token = None + self._claims = {} + + browser_loop = self._browser_loop + + try: + if browser_loop is not None and browser_loop.is_running(): + future = asyncio.run_coroutine_threadsafe( + self._close_browser_resources_async(), + browser_loop, + ) + await asyncio.wrap_future(future) + else: + await self._close_browser_resources_async() + finally: + await self._stop_browser_thread_async() + + async def get_claims_async(self) -> dict[str, Any]: + """Return the claims extracted from the current token.""" + return dict(self._claims) + + def _has_fresh_token(self) -> bool: + """Return whether the current token is outside of the expiry buffer.""" + if self._access_token is None: + return False + + expires_at: object = self._claims.get("exp") + if not isinstance(expires_at, (int, float)): + return False + + return expires_at - time.time() > self._expiry_buffer_seconds + + async def _capture_access_token_async(self) -> str: + """ + Capture a token on the authenticator-owned browser event loop. + + Returns: + The captured access token as a string. + """ + return await self._run_on_browser_thread_async( + operation=self._capture_access_token_on_browser_loop_async, + ) + + async def _run_on_browser_thread_async( + self, + *, + operation: Callable[[], Coroutine[Any, Any, T]], + ) -> T: + """ + Run a browser operation on the retained browser event loop. + + Args: + operation: An asynchronous callable representing the browser operation to run. + + Returns: + The result of the browser operation. + + Raises: + RuntimeError: If the browser event loop failed to start. + """ + await self._ensure_browser_thread_started_async() + + browser_loop = self._browser_loop + if browser_loop is None: + raise RuntimeError("Browser event loop failed to start.") + + future = asyncio.run_coroutine_threadsafe(operation(), browser_loop) + return await asyncio.wrap_future(future) + + async def _ensure_browser_thread_started_async(self) -> None: + """ + Start the retained browser event-loop thread when needed. + + Raises: + RuntimeError: If the browser event loop failed to start. + """ + thread = self._browser_thread + if thread is not None and thread.is_alive(): + return + + self._browser_loop_started.clear() + self._browser_loop_start_error = None + thread = threading.Thread( + target=self._run_browser_event_loop, + name="pyrit-copilot-browser", + daemon=True, + ) + self._browser_thread = thread + thread.start() + + await asyncio.to_thread(self._browser_loop_started.wait) + + if self._browser_loop_start_error is not None: + raise RuntimeError("Browser event loop failed to start.") from self._browser_loop_start_error + + def _run_browser_event_loop(self) -> None: + """Run the browser event loop on its dedicated thread.""" + try: + loop = asyncio.ProactorEventLoop() if sys.platform == "win32" else asyncio.new_event_loop() + asyncio.set_event_loop(loop) + self._browser_loop = loop + self._browser_loop_started.set() + loop.run_forever() + loop.close() + except Exception as error: + self._browser_loop_start_error = error + finally: + self._browser_loop_started.set() + self._browser_loop = None + + async def _stop_browser_thread_async(self) -> None: + """Stop and join the retained browser thread.""" + loop = self._browser_loop + thread = self._browser_thread + + if loop is not None and loop.is_running(): + loop.call_soon_threadsafe(loop.stop) + + if thread is not None and thread.is_alive(): + await asyncio.to_thread(thread.join) + + self._browser_loop = None + self._browser_thread = None + + async def _capture_access_token_on_browser_loop_async(self) -> str: + """ + Capture a Copilot access token from a persistent Edge session. + + Returns: + The captured access token as a string. + """ + token_future: asyncio.Future[str] = asyncio.get_running_loop().create_future() + page = await self._ensure_browser_session_async() + + def handle_websocket(websocket: Any) -> None: + self._handle_websocket_url( + websocket_url=str(websocket.url), + token_future=token_future, + ) + + page.on("websocket", handle_websocket) + + try: + await page.goto(self._copilot_url) + access_token = await self._wait_for_token_async( + token_future=token_future, + ) + + if not self._headless: + await self._minimize_browser_window_async(page=page) + + return access_token + except BaseException: + await self._close_browser_resources_async() + raise + finally: + page.remove_listener("websocket", handle_websocket) + + async def _capture_and_store_token_async(self) -> str: + """ + Capture, validate and store a new Copilot access token. + + Returns: + The newly captured access token as a string. + + Raises: + ValueError: If captured token claims are missing, invalid, or within the expiry buffer. + """ + captured_token = await self._capture_access_token_async() + claims = jwt.decode( + captured_token, + algorithms=["RS256"], + options={"verify_signature": False}, + ) + + required_claims = {"tid", "oid", "exp"} + missing_claims = required_claims - claims.keys() + if missing_claims: + raise ValueError(f"Missing required claim: {', '.join(missing_claims)}") + + tenant_id = claims["tid"] + object_id = claims["oid"] + expires_at = claims["exp"] + + if not isinstance(tenant_id, str) or not tenant_id.strip(): + raise ValueError("Captured Copilot token has an invalid tid claim.") + + if not isinstance(object_id, str) or not object_id.strip(): + raise ValueError("Captured Copilot token has an invalid oid claim.") + + if isinstance(expires_at, bool) or not isinstance(expires_at, (int, float)): + raise ValueError("Captured Copilot token has an invalid exp claim.") + + if expires_at - time.time() <= self._expiry_buffer_seconds: + raise ValueError("Captured Copilot token is within the expiry buffer.") + + self._claims = claims + self._access_token = captured_token + + return captured_token + + def _extract_access_token_from_websocket_url(self, *, websocket_url: str) -> str | None: + """ + Extract the access token from a Copilot Chathub websocket URL. + + Args: + websocket_url: The full websocket URL containing the access token. + + Returns: + The access token if present, otherwise None. + """ + parsed_url = urlparse(websocket_url) + expected_url = urlparse(self._websocket_base_url) + + if ( + parsed_url.scheme != expected_url.scheme + or parsed_url.hostname != expected_url.hostname + or not parsed_url.path.startswith(f"{expected_url.path}/") + ): + return None + + tokens = parse_qs(parsed_url.query).get("access_token") + return tokens[0] if tokens else None + + def _handle_websocket_url( + self, + *, + websocket_url: str, + token_future: asyncio.Future[str], + ) -> None: + """Resolve the token future when a valid Copilot url is observed.""" + if token_future.done(): + return + + access_token = self._extract_access_token_from_websocket_url(websocket_url=websocket_url) + if access_token: + token_future.set_result(access_token) + + async def _wait_for_token_async( + self, + *, + token_future: asyncio.Future[str], + ) -> str: + """ + Wait for the token to be captured or raise a timeout. + + Args: + token_future: The future that will be resolved with the access token. + + Returns: + The captured access token. + + Raises: + TimeoutError: If the token is not captured within the specified timeout. + """ + try: + return await asyncio.wait_for( + token_future, + timeout=self._token_capture_timeout_seconds, + ) + except asyncio.TimeoutError: + raise TimeoutError( + "Timed out waiting for access token capture. Complete sign-in in the opened browser and try again." + ) from None + + @staticmethod + def _create_playwright_context_manager() -> Any: + """ + Create Playwright's asynchronous context manager. + + Returns: + An instance of Playwright's asynchronous context manager. + + Raises: + RuntimeError: If Playwright is not installed. + """ + try: + from playwright.async_api import async_playwright # type: ignore[ty:unresolved-import] + except ImportError: + raise RuntimeError( + "Playwright is required for browser-session authentication. " + "Install the PyRIT playwright extra and run 'playwright install msedge'." + ) from None + + return async_playwright() + + async def _ensure_browser_session_async(self) -> Any: + """ + Ensure that the browser session is initialized and return the page instance. + + Returns: + The initialized page instance. + """ + if self._page is not None: + return self._page + + resources = AsyncExitStack() + try: + playwright = await resources.enter_async_context(self._create_playwright_context_manager()) + browser_context = await playwright.chromium.launch_persistent_context( + user_data_dir=str(self._profile_path), + channel=self._browser_channel, + headless=self._headless, + ) + resources.push_async_callback(browser_context.close) + + pages = browser_context.pages + page = pages[0] if pages else await browser_context.new_page() + except BaseException: + await resources.aclose() + raise + + self._browser_resources = resources + self._browser_context = browser_context + self._page = page + return page + + async def _minimize_browser_window_async(self, *, page: Any) -> None: + """ + Minimize the browser window using the Chrome DevTools Protocol (CDP). + + Args: + page: The page instance representing the browser tab. + + Raises: + RuntimeError: If minimizing the browser window fails. + """ + cdp_session = await page.context.new_cdp_session(page) + + try: + window_info = await cdp_session.send("Browser.getWindowForTarget") + window_id = window_info["windowId"] + + if not isinstance(window_id, int): + raise RuntimeError("Edge did not return a valid browser window ID.") + + await cdp_session.send( + "Browser.setWindowBounds", + {"windowId": window_id, "bounds": {"windowState": "minimized"}}, + ) + finally: + await cdp_session.detach() + + async def _close_browser_resources_async(self) -> None: + """Close browser resources on their owning event loop.""" + resources = self._browser_resources + + self._browser_resources = None + self._browser_context = None + self._page = None + + if resources is not None: + await resources.aclose() diff --git a/pyrit/auth/manual_copilot_authenticator.py b/pyrit/auth/manual_copilot_authenticator.py index 3fc8e4070f..99d10793a3 100644 --- a/pyrit/auth/manual_copilot_authenticator.py +++ b/pyrit/auth/manual_copilot_authenticator.py @@ -25,7 +25,7 @@ class ManualCopilotAuthenticator(Authenticator): 3. Go to the Network tab. 4. Filter by "Socket" connections or search for "Chathub". 5. Start typing in the chat to initiate a WebSocket connection. - 6. Look for the latest WebSocket connection to ``substrate.office.com/m365Copilot/Chathub``. + 6. Look for the latest WebSocket connection to ``substrate.svc.cloud.microsoft/m365Copilot/Chathub``. 7. You may find the ``access_token`` in the request URL or in the request payload. Note: diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index fcd8113f36..107e5a989a 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -13,7 +13,11 @@ import websockets from websockets.exceptions import InvalidStatus -from pyrit.auth import CopilotAuthenticator, ManualCopilotAuthenticator +from pyrit.auth import ( + BrowserSessionCopilotAuthenticator, + CopilotAuthenticator, + ManualCopilotAuthenticator, +) from pyrit.common import get_mime_type from pyrit.exceptions import ( EmptyResponseException, @@ -53,13 +57,18 @@ class WebSocketCopilotTarget(PromptTarget): A WebSocket-based prompt target for integrating with Microsoft Copilot. This class facilitates communication with Microsoft Copilot over a WebSocket connection. - Authentication can be handled in two ways: + Authentication can be handled in three ways: + + 1. **Credential automation (default)**: ``CopilotAuthenticator`` reads + ``COPILOT_USERNAME`` and ``COPILOT_PASSWORD``, automates the login form with + Playwright, and caches the captured token. - 1. **Automated (default)**: Via ``CopilotAuthenticator``, which uses Playwright to automate - browser login and obtain the required access tokens. Requires ``COPILOT_USERNAME`` and - ``COPILOT_PASSWORD`` environment variables as well as Playwright installed. + 2. **Interactive browser session**: ``BrowserSessionCopilotAuthenticator`` opens + a persistent Edge profile for user-driven sign-in, including MFA and + Conditional Access. It keeps captured tokens in memory and refreshes them + through the retained browser session. - 2. **Manual**: Via ``ManualCopilotAuthenticator``, which accepts a pre-obtained access token. + 3. **Manual**: Via ``ManualCopilotAuthenticator``, which accepts a pre-obtained access token. This is useful for situations where browser automation is not possible. Once authenticated, the target supports multi-turn conversations through server-side @@ -95,11 +104,13 @@ class WebSocketCopilotTarget(PromptTarget): def __init__( self, *, - websocket_base_url: str = "wss://substrate.office.com/m365Copilot/Chathub", + websocket_base_url: str = "wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub", max_requests_per_minute: int | None = None, model_name: str = "copilot", response_timeout_seconds: int = RESPONSE_TIMEOUT_SECONDS, - authenticator: CopilotAuthenticator | ManualCopilotAuthenticator | None = None, + authenticator: ( + BrowserSessionCopilotAuthenticator | CopilotAuthenticator | ManualCopilotAuthenticator | None + ) = None, custom_configuration: TargetConfiguration | None = None, ) -> None: """ @@ -107,13 +118,14 @@ def __init__( Args: websocket_base_url (str): Base URL for the Copilot WebSocket endpoint. - Defaults to ``wss://substrate.office.com/m365Copilot/Chathub``. + Defaults to ``wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub``. max_requests_per_minute (int | None): Maximum number of requests per minute. model_name (str): The model name. Defaults to "copilot". response_timeout_seconds (int): Timeout for receiving responses in seconds. Defaults to 60s. - authenticator (CopilotAuthenticator | ManualCopilotAuthenticator | None): Authenticator - instance. Supports both ``CopilotAuthenticator`` and ``ManualCopilotAuthenticator``. - If None, a new ``CopilotAuthenticator`` instance will be created with default settings. + authenticator (Authenticator | None): Copilot authentication strategy. Supported + implementations are ``BrowserSessionCopilotAuthenticator``, + ``CopilotAuthenticator``, and ``ManualCopilotAuthenticator``. If None, a new + ``CopilotAuthenticator`` is created. custom_configuration (TargetConfiguration, Optional): Override the default configuration for this target instance. Defaults to None. diff --git a/tests/unit/auth/test_browser_session_copilot_authenticator.py b/tests/unit/auth/test_browser_session_copilot_authenticator.py new file mode 100644 index 0000000000..b708695585 --- /dev/null +++ b/tests/unit/auth/test_browser_session_copilot_authenticator.py @@ -0,0 +1,754 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +import threading +import time +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, call, patch + +import jwt +import pytest + +from pyrit.auth import ( + BrowserSessionCopilotAuthenticator, +) +from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH + +_TEST_JWT_KEY = "a" * 32 + + +def _make_token( + *, + claims: dict[str, object] | None = None, +) -> str: + """Create a JWT token with a short expiry for testing.""" + token_claims = ( + claims + if claims is not None + else { + "tid": "tenant_id", + "oid": "object_id", + "exp": int(time.time()) + 3600, + } + ) + return jwt.encode( + token_claims, + key=_TEST_JWT_KEY, + algorithm="HS256", + ) + + +def test_init_uses_custom_profile_path(tmp_path: Path) -> None: + + profile_path = tmp_path / "copilot_profile" + + authenticator = BrowserSessionCopilotAuthenticator( + profile_path=profile_path, + ) + + assert authenticator.profile_path == profile_path + + +def test_init_uses_default_profile_path() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + assert authenticator.profile_path == (CONFIGURATION_DIRECTORY_PATH / "copilot_browser_profiles" / "default") + + +@pytest.mark.parametrize("timeout", [0, -1]) +def test_init_rejects_non_positive_capture_timeout(timeout: int) -> None: + with pytest.raises(ValueError, match="token_capture_timeout_seconds must be a positive integer."): + BrowserSessionCopilotAuthenticator( + token_capture_timeout_seconds=timeout, + ) + + +@pytest.mark.parametrize("expiry_buffer", [0, -1]) +def test_init_rejects_non_positive_expiry_buffer(expiry_buffer: int) -> None: + with pytest.raises(ValueError, match="expiry_buffer_seconds must be a positive integer."): + BrowserSessionCopilotAuthenticator( + expiry_buffer_seconds=expiry_buffer, + ) + + +async def test_get_token_async_captures_token_when_missing() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + captured_token = _make_token() + + with patch.object( + BrowserSessionCopilotAuthenticator, + "_capture_access_token_async", + new=AsyncMock(return_value=captured_token), + create=True, + ) as capture: + result = await authenticator.get_token_async() + + assert result == captured_token + capture.assert_awaited_once() + + +async def test_get_claims_async_returns_captured_token_claims() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + token = _make_token() + + with patch.object( + BrowserSessionCopilotAuthenticator, + "_capture_access_token_async", + new=AsyncMock(return_value=token), + ): + await authenticator.get_token_async() + + claims = await authenticator.get_claims_async() + + assert claims["tid"] == "tenant_id" + assert claims["oid"] == "object_id" + + +@pytest.mark.parametrize("missing_claim", ["tid", "oid", "exp"]) +async def test_get_token_async_rejects_missing_required_claims(missing_claim: str) -> None: + claims = { + "tid": "tenant_id", + "oid": "object_id", + "exp": int(time.time()) + 3600, + } + claims.pop(missing_claim) + + token = _make_token(claims=claims) + authenticator = BrowserSessionCopilotAuthenticator() + + with patch.object( + BrowserSessionCopilotAuthenticator, + "_capture_access_token_async", + new=AsyncMock(return_value=token), + ): + with pytest.raises(ValueError, match=f"Missing required claim: {missing_claim}"): + await authenticator.get_token_async() + + +async def test_get_token_async_recaptures_token_within_expiry_buffer() -> None: + authenticator = BrowserSessionCopilotAuthenticator( + expiry_buffer_seconds=300, + ) + initial_token = _make_token( + claims={ + "tid": "tenant_id", + "oid": "object_id", + "exp": 1000, + } + ) + refreshed_token = _make_token( + claims={ + "tid": "tenant_id", + "oid": "object_id", + "exp": 2000, + } + ) + + capture = AsyncMock(side_effect=[initial_token, refreshed_token]) + + with ( + patch.object( + BrowserSessionCopilotAuthenticator, + "_capture_access_token_async", + new=capture, + ), + patch( + "pyrit.auth.browser_session_copilot_authenticator.time.time", + return_value=0, + ) as current_time, + ): + assert await authenticator.get_token_async() == initial_token + + current_time.return_value = 750 + + assert await authenticator.get_token_async() == refreshed_token + + assert capture.await_count == 2 + + +async def test_get_token_async_serializes_concurrent_capture() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + token = _make_token() + capture_started = asyncio.Event() + release_capture = asyncio.Event() + + async def capture_token_async() -> str: + capture_started.set() + await release_capture.wait() + return token + + with patch.object( + BrowserSessionCopilotAuthenticator, + "_capture_access_token_async", + side_effect=capture_token_async, + ) as capture: + first_request = asyncio.create_task(authenticator.get_token_async()) + await capture_started.wait() + + second_request = asyncio.create_task(authenticator.get_token_async()) + release_capture.set() + + results = await asyncio.gather(first_request, second_request) + + assert results == [token, token] + capture.assert_awaited_once() + + +async def test_get_token_async_captures_independently_per_instance( + tmp_path: Path, +) -> None: + first = BrowserSessionCopilotAuthenticator(profile_path=tmp_path / "persona_one") + second = BrowserSessionCopilotAuthenticator(profile_path=tmp_path / "persona_two") + + first_token = _make_token( + claims={ + "tid": "tenant_id", + "oid": "object_id_one", + "exp": int(time.time()) + 3600, + } + ) + second_token = _make_token( + claims={ + "tid": "tenant_id", + "oid": "object_id_two", + "exp": int(time.time()) + 3600, + } + ) + + with patch.object( + BrowserSessionCopilotAuthenticator, + "_capture_access_token_async", + new=AsyncMock(side_effect=[first_token, second_token]), + ) as capture: + results = await asyncio.gather( + first.get_token_async(), + second.get_token_async(), + ) + + assert results == [first_token, second_token] + assert capture.await_count == 2 + + +async def test_refresh_token_async_forces_new_capture() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + initial_token = _make_token() + refreshed_token = _make_token( + claims={ + "tid": "tenant_id", + "oid": "object_id", + "exp": int(time.time()) + 7200, + } + ) + capture = AsyncMock(side_effect=[initial_token, refreshed_token]) + + with patch.object( + BrowserSessionCopilotAuthenticator, + "_capture_access_token_async", + new=capture, + ): + assert await authenticator.get_token_async() == initial_token + assert await authenticator.refresh_token_async() == refreshed_token + + assert capture.await_count == 2 + + +def test_extract_access_token_from_chathub_url() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + websocket_url = "wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub/user@tenant?access_token=test-token" + + result = authenticator._extract_access_token_from_websocket_url(websocket_url=websocket_url) + + assert result == "test-token" + + +@pytest.mark.parametrize( + "websocket_url", + [ + "ws://substrate.svc.cloud.microsoft/m365Copilot/Chathub/user@tenant?access_token=test-token", + "wss://evil.example/m365Copilot/Chathub/user@tenant?access_token=test-token", + "wss://substrate.svc.cloud.microsoft/other/path?access_token=test-token", + "wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub/user@tenant", + ], +) +def test_extract_access_token_rejects_unexpected_url( + websocket_url: str, +) -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + result = authenticator._extract_access_token_from_websocket_url(websocket_url=websocket_url) + assert result is None + + +async def test_handle_websocket_url_resolves_token_future() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + token_future = asyncio.get_running_loop().create_future() + websocket_url = "wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub/user@tenant?access_token=test-token" + + result = authenticator._handle_websocket_url( + websocket_url=websocket_url, + token_future=token_future, + ) + + assert await token_future == "test-token" + + +async def test_wait_for_token_async_raises_clear_timeout() -> None: + authenticator = BrowserSessionCopilotAuthenticator( + token_capture_timeout_seconds=1, + ) + token_future: asyncio.Future[str] = asyncio.get_running_loop().create_future() + + with pytest.raises( + TimeoutError, + match="Timed out waiting for access token capture.", + ): + await authenticator._wait_for_token_async(token_future=token_future) + + +def test_init_uses_browser_capture_defaults() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + assert authenticator._browser_channel == "msedge" + assert authenticator._headless is False + assert authenticator._copilot_url == "https://m365.cloud.microsoft/chat" + + +async def test_capture_access_token_async_keeps_context_until_closed( + tmp_path: Path, +) -> None: + profile_path = tmp_path / "copilot-profile" + authenticator = BrowserSessionCopilotAuthenticator( + profile_path=profile_path, + ) + + page = MagicMock() + page.goto = AsyncMock() + + browser_context = MagicMock() + browser_context.pages = [] + browser_context.new_page = AsyncMock(return_value=page) + browser_context.close = AsyncMock() + + playwright = MagicMock() + playwright.chromium.launch_persistent_context = AsyncMock(return_value=browser_context) + + playwright_manager = AsyncMock() + playwright_manager.__aenter__.return_value = playwright + playwright_manager.__aexit__.return_value = None + + with ( + patch.object( + authenticator, + "_create_playwright_context_manager", + return_value=playwright_manager, + create=True, + ), + patch.object( + authenticator, + "_wait_for_token_async", + new=AsyncMock(return_value="test-token"), + ), + patch.object( + authenticator, + "_minimize_browser_window_async", + new=AsyncMock(), + ) as minimize, + ): + result = await authenticator._capture_access_token_async() + + assert result == "test-token" + playwright.chromium.launch_persistent_context.assert_awaited_once_with( + user_data_dir=str(profile_path), + channel="msedge", + headless=False, + ) + browser_context.new_page.assert_awaited_once() + page.on.assert_called_once() + page.goto.assert_awaited_once_with("https://m365.cloud.microsoft/chat") + minimize.assert_awaited_once_with(page=page) + browser_context.close.assert_not_awaited() + playwright_manager.__aexit__.assert_not_awaited() + + await authenticator.close_async() + + browser_context.close.assert_awaited_once() + playwright_manager.__aexit__.assert_awaited_once() + + +async def test_capture_access_token_async_closes_context_on_navigation_error( + tmp_path: Path, +) -> None: + authenticator = BrowserSessionCopilotAuthenticator( + profile_path=tmp_path / "copilot-profile", + ) + + page = MagicMock() + page.goto = AsyncMock(side_effect=RuntimeError("navigation failed")) + + browser_context = MagicMock() + browser_context.new_page = AsyncMock(return_value=page) + browser_context.close = AsyncMock() + browser_context.pages = [] + + playwright = MagicMock() + playwright.chromium.launch_persistent_context = AsyncMock(return_value=browser_context) + + playwright_manager = AsyncMock() + playwright_manager.__aenter__.return_value = playwright + playwright_manager.__aexit__.return_value = None + + with ( + patch.object( + authenticator, + "_create_playwright_context_manager", + return_value=playwright_manager, + ), + pytest.raises(RuntimeError, match="navigation failed"), + ): + await authenticator._capture_access_token_async() + + browser_context.close.assert_awaited_once() + + +def test_auth_package_exports_browser_session_authenticator() -> None: + from pyrit.auth import BrowserSessionCopilotAuthenticator as ExportedAuthenticator + + assert ExportedAuthenticator is BrowserSessionCopilotAuthenticator + + +async def test_capture_access_token_async_reuses_existing_page() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + page = MagicMock() + page.on = MagicMock() + page.remove_listener = MagicMock() + page.goto = AsyncMock() + + authenticator._page = page + + with ( + patch.object( + authenticator, + "_wait_for_token_async", + new=AsyncMock(return_value="refreshed_token"), + ), + patch.object( + authenticator, + "_minimize_browser_window_async", + new=AsyncMock(), + ) as minimize, + ): + result = await authenticator._capture_access_token_async() + + assert result == "refreshed_token" + page.goto.assert_awaited_once_with("https://m365.cloud.microsoft/chat") + page.on.assert_called_once() + page.remove_listener.assert_called_once() + minimize.assert_awaited_once_with(page=page) + + +async def test_minimize_browser_window_async() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + cdp_session = MagicMock() + cdp_session.send = AsyncMock( + side_effect=[ + {"windowId": 42}, + {}, + ] + ) + cdp_session.detach = AsyncMock() + + page = MagicMock() + page.context.new_cdp_session = AsyncMock(return_value=cdp_session) + + await authenticator._minimize_browser_window_async(page=page) + + assert cdp_session.send.await_args_list == [ + call("Browser.getWindowForTarget"), + call("Browser.setWindowBounds", {"windowId": 42, "bounds": {"windowState": "minimized"}}), + ] + cdp_session.detach.assert_awaited_once() + + +async def test_close_async_is_idempotent() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + resources = AsyncMock() + authenticator._browser_resources = resources + + await authenticator.close_async() + await authenticator.close_async() + + resources.aclose.assert_awaited_once() + assert authenticator._browser_resources is None + assert authenticator._browser_context is None + assert authenticator._page is None + + +async def test_async_context_manager_returns_authenticator() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + async with authenticator as entered: + assert entered is authenticator + + +async def test_async_context_manager_closes_on_exit() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + with patch.object( + authenticator, + "close_async", + new=AsyncMock(), + ) as close: + async with authenticator: + pass + + close.assert_awaited_once() + + +async def test_close_async_clears_token_state() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + token = _make_token() + + with patch.object( + authenticator, + "_capture_access_token_async", + new=AsyncMock(return_value=token), + ): + await authenticator.get_token_async() + + await authenticator.close_async() + + assert authenticator._access_token is None + assert await authenticator.get_claims_async() == {} + + +async def test_capture_access_token_async_always_delegates_to_browser_thread() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + with ( + patch.object( + authenticator, + "_ensure_browser_session_async", + new=AsyncMock(side_effect=AssertionError("Browser must not launch in this unit test")), + ), + patch.object( + authenticator, + "_run_on_browser_thread_async", + new=AsyncMock(return_value="test-token"), + create=True, + ) as run_on_thread, + ): + result = await authenticator._capture_access_token_async() + + assert result == "test-token" + run_on_thread.assert_awaited_once() + + +async def test_run_on_browser_thread_async_reuses_thread_and_loop() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + async def identify_execution_context_async() -> tuple[int, int]: + return threading.get_ident(), id(asyncio.get_running_loop()) + + try: + first_context = await authenticator._run_on_browser_thread_async( + operation=identify_execution_context_async, + ) + second_context = await authenticator._run_on_browser_thread_async( + operation=identify_execution_context_async, + ) + finally: + await authenticator.close_async() + + assert first_context == second_context + assert first_context[0] != threading.get_ident() + + +async def test_close_async_closes_resources_on_browser_thread() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + resources = MagicMock() + cleanup_thread_id: int | None = None + + async def close_resources_async() -> None: + nonlocal cleanup_thread_id + cleanup_thread_id = threading.get_ident() + + resources.aclose = AsyncMock(side_effect=close_resources_async) + + async def install_resources_async() -> int: + authenticator._browser_resources = resources + return threading.get_ident() + + owner_thread_id = await authenticator._run_on_browser_thread_async( + operation=install_resources_async, + ) + + await authenticator.close_async() + + assert cleanup_thread_id == owner_thread_id + + +def test_extract_access_token_accepts_trailing_slash_in_base_url() -> None: + authenticator = BrowserSessionCopilotAuthenticator( + websocket_base_url=("wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub/"), + ) + + result = authenticator._extract_access_token_from_websocket_url( + websocket_url=("wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub/user@tenant?access_token=test-token"), + ) + + assert result == "test-token" + + +async def test_get_token_async_rejects_new_token_within_expiry_buffer() -> None: + authenticator = BrowserSessionCopilotAuthenticator( + expiry_buffer_seconds=300, + ) + token = _make_token( + claims={ + "tid": "tenant_id", + "oid": "object_id", + "exp": 1200, + } + ) + + with ( + patch.object( + authenticator, + "_capture_access_token_async", + new=AsyncMock(return_value=token), + ), + patch( + "pyrit.auth.browser_session_copilot_authenticator.time.time", + return_value=1000, + ), + pytest.raises(ValueError, match="within the expiry buffer"), + ): + await authenticator.get_token_async() + + assert authenticator._access_token is None + assert await authenticator.get_claims_async() == {} + + +@pytest.mark.parametrize( + ("claim", "value", "message"), + [ + ("tid", "", "invalid tid claim"), + ("oid", " ", "invalid oid claim"), + ("exp", True, "invalid exp claim"), + ("exp", "tomorrow", "invalid exp claim"), + ], +) +async def test_get_token_async_rejects_invalid_required_claim( + claim: str, + value: object, + message: str, +) -> None: + claims: dict[str, object] = { + "tid": "tenant_id", + "oid": "object_id", + "exp": int(time.time()) + 3600, + } + claims[claim] = value + token = _make_token(claims=claims) + authenticator = BrowserSessionCopilotAuthenticator() + + with ( + patch.object( + authenticator, + "_capture_access_token_async", + new=AsyncMock(return_value=token), + ), + pytest.raises(ValueError, match=message), + ): + await authenticator.get_token_async() + + assert authenticator._access_token is None + + +async def test_close_async_waits_for_active_token_operation() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + await authenticator._token_fetch_lock.acquire() + close_task = asyncio.create_task(authenticator.close_async()) + + await asyncio.sleep(0) + + assert not close_task.done() + + authenticator._token_fetch_lock.release() + await close_task + + +async def test_ensure_browser_session_async_cleans_up_on_cancellation() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + playwright = MagicMock() + playwright.chromium.launch_persistent_context = AsyncMock( + side_effect=asyncio.CancelledError, + ) + + playwright_manager = AsyncMock() + playwright_manager.__aenter__.return_value = playwright + playwright_manager.__aexit__.return_value = None + + with ( + patch.object( + authenticator, + "_create_playwright_context_manager", + return_value=playwright_manager, + ), + pytest.raises(asyncio.CancelledError), + ): + await authenticator._ensure_browser_session_async() + + playwright_manager.__aexit__.assert_awaited_once() + + +async def test_capture_on_browser_loop_cleans_up_on_cancellation() -> None: + authenticator = BrowserSessionCopilotAuthenticator() + + page = MagicMock() + page.on = MagicMock() + page.remove_listener = MagicMock() + page.goto = AsyncMock(side_effect=asyncio.CancelledError) + + with ( + patch.object( + authenticator, + "_ensure_browser_session_async", + new=AsyncMock(return_value=page), + ), + patch.object( + authenticator, + "_close_browser_resources_async", + new=AsyncMock(), + ) as close_resources, + pytest.raises(asyncio.CancelledError), + ): + await authenticator._capture_access_token_on_browser_loop_async() + + close_resources.assert_awaited_once() + page.remove_listener.assert_called_once() + + +@pytest.mark.parametrize( + "websocket_base_url", + [ + "ws://substrate.svc.cloud.microsoft/m365Copilot/Chathub", + "https://substrate.svc.cloud.microsoft/m365Copilot/Chathub", + "wss:///m365Copilot/Chathub", + "wss://substrate.svc.cloud.microsoft", + ], +) +def test_init_rejects_invalid_websocket_base_url( + websocket_base_url: str, +) -> None: + with pytest.raises( + ValueError, + match="websocket_base_url must be a valid wss URL with a path", + ): + BrowserSessionCopilotAuthenticator( + websocket_base_url=websocket_base_url, + ) diff --git a/tests/unit/prompt_target/target/test_websocket_copilot_target.py b/tests/unit/prompt_target/target/test_websocket_copilot_target.py index 936def927a..47e52be345 100644 --- a/tests/unit/prompt_target/target/test_websocket_copilot_target.py +++ b/tests/unit/prompt_target/target/test_websocket_copilot_target.py @@ -2,12 +2,13 @@ # Licensed under the MIT license. import asyncio +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import jwt import pytest -from pyrit.auth import CopilotAuthenticator +from pyrit.auth import BrowserSessionCopilotAuthenticator, CopilotAuthenticator from pyrit.models import Message, MessagePiece from pyrit.prompt_target import WebSocketCopilotTarget from pyrit.prompt_target.websocket_copilot_target import CopilotMessageType @@ -139,7 +140,7 @@ def test_init_with_default_parameters(self): assert target._authenticator == mock_auth_instance assert target._response_timeout_seconds == WebSocketCopilotTarget.RESPONSE_TIMEOUT_SECONDS assert target._model_name == "copilot" - assert target._endpoint == "wss://substrate.office.com/m365Copilot/Chathub" + assert target._endpoint == "wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub" assert target._verbose is False assert target._max_requests_per_minute is None @@ -161,6 +162,15 @@ def test_init_with_invalid_response_timeout(self, mock_authenticator): with pytest.raises(ValueError, match="response_timeout_seconds must be a positive integer."): WebSocketCopilotTarget(authenticator=mock_authenticator, response_timeout_seconds=invalid_timeout) + def test_init_with_browser_session_authenticator(self, tmp_path: Path) -> None: + authenticator = BrowserSessionCopilotAuthenticator( + profile_path=tmp_path / "copilot-profile", + ) + + target = WebSocketCopilotTarget(authenticator=authenticator) + + assert target._authenticator is authenticator + @pytest.mark.usefixtures("patch_central_database") class TestDictToWebsocket: @@ -252,7 +262,7 @@ async def test_build_websocket_url_with_valid_token(self, mock_authenticator, mo ) expected_token = await mock_authenticator.get_token_async() - assert url.startswith("wss://substrate.office.com/m365Copilot/Chathub/test_object_id@test_tenant_id?") + assert url.startswith("wss://substrate.svc.cloud.microsoft/m365Copilot/Chathub/test_object_id@test_tenant_id?") assert f"X-SessionId={session_id}" in url assert f"ConversationId={copilot_conversation_id}" in url assert f"access_token={expected_token}" in url