From 9823509ba8f8a39f6a340f9880879657617d3d3e Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Mon, 22 Dec 2025 20:43:09 +0100 Subject: [PATCH 01/26] initial, minimal working version of WebSocketCopilotTarget --- doc/api.rst | 1 + pyrit/prompt_target/__init__.py | 2 + .../prompt_target/websocket_copilot_target.py | 297 ++++++++++++++++++ websocket_copilot_simple_example.py | 31 ++ 4 files changed, 331 insertions(+) create mode 100644 pyrit/prompt_target/websocket_copilot_target.py create mode 100644 websocket_copilot_simple_example.py diff --git a/doc/api.rst b/doc/api.rst index 1b9bdc775..cd812de72 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -500,6 +500,7 @@ API Reference PromptTarget RealtimeTarget TextTarget + WebSocketCopilotTarget :py:mod:`pyrit.score` ===================== diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index cdbbdb0ff..eee3ff6cb 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -37,6 +37,7 @@ from pyrit.prompt_target.playwright_copilot_target import CopilotType, PlaywrightCopilotTarget from pyrit.prompt_target.prompt_shield_target import PromptShieldTarget from pyrit.prompt_target.text_target import TextTarget +from pyrit.prompt_target.websocket_copilot_target import WebSocketCopilotTarget __all__ = [ "AzureBlobStorageTarget", @@ -66,4 +67,5 @@ "PromptTarget", "RealtimeTarget", "TextTarget", + "WebSocketCopilotTarget", ] diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py new file mode 100644 index 000000000..f15994f56 --- /dev/null +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -0,0 +1,297 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json +import logging +import os +import uuid +from enum import Enum +from typing import Optional + +import websockets + +from pyrit.models import Message, construct_response_from_request +from pyrit.prompt_target import PromptTarget, limit_requests_per_minute + +logger = logging.getLogger(__name__) + +""" +Useful links: +https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py +https://labs.zenity.io/p/access-copilot-m365-terminal +""" + + +class CopilotMessageType(Enum): + """Enumeration for Copilot WebSocket message types.""" + + UNKNOWN = -1 + NEXT_DATA_FRAME = 1 # streaming Copilot responses + LAST_DATA_FRAME = 2 # the last data frame with final content + USER_PROMPT = 4 + PING = 6 + + +class WebSocketCopilotTarget(PromptTarget): + """ + A WebSocket-based prompt target for Microsoft Copilot integration. + + This target enables communication with Microsoft Copilot through a WebSocket connection. + Currently, authentication requires manually extracting a WebSocket URL from an active browser session. + In the future, more flexible authentication mechanisms will be added. + + To obtain the WebSocket URL: + 1. Ensure you are logged into Microsoft 365 with access to Copilot + 2. Navigate to https://m365.cloud.microsoft/chat or open Copilot in https://teams.microsoft.com/v2 + 3. Open browser developer tools and switch to the Network tab + 4. Begin typing or send a message to Copilot to establish the WebSocket connection + 5. Search the network requests for "chathub", "conversation", or "access_token" + 6. Identify the WebSocket connection (look for WS protocol) and copy its full URL + + Warning: + All target instances using the same `WEBSOCKET_URL` will share a single conversation session. + Only works with licensed Microsoft 365 Copilot. The free Copilot version is not compatible. + """ + + # TODO: add more flexible auth, use puppeteer? https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py#L248 + # TODO: add useful message for: "Error during WebSocket communication: server rejected WebSocket connection: HTTP 401" + + SUPPORTED_DATA_TYPES = {"text"} # TODO: support more types? + + # TODO: implement timeouts and retries + MAX_WAIT_TIME_SECONDS: int = 300 + POLL_INTERVAL_MS: int = 2000 + + def __init__( + self, + *, + verbose: bool = False, + max_requests_per_minute: Optional[int] = None, + model_name: str = "copilot", + ) -> None: + """ + Initialize the WebSocketCopilotTarget. + + Args: + verbose (bool): Enable verbose logging. Defaults to False. + max_requests_per_minute (int, Optional): Maximum number of requests per minute. + model_name (str): The model name. Defaults to "copilot". + + Raises: + ValueError: If WebSocket URL is not provided as env variable. + """ + self._websocket_url = os.getenv("WEBSOCKET_URL") + if not self._websocket_url: + raise ValueError("WebSocket URL must be provided through the WEBSOCKET_URL environment variable") + + if not "ConversationId=" in self._websocket_url: + raise ValueError("`ConversationId` parameter not found in URL.") + self._conversation_id = self._websocket_url.split("ConversationId=")[1].split("&")[0] + + if not "X-SessionId=" in self._websocket_url: + raise ValueError("`X-SessionId` parameter not found in URL.") + self._session_id = self._websocket_url.split("X-SessionId=")[1].split("&")[0] + + super().__init__( + verbose=verbose, + max_requests_per_minute=max_requests_per_minute, + endpoint=self._websocket_url.split("?")[0], # wss://substrate.office.com/m365Copilot/Chathub/... + model_name=model_name, + ) + + if self._verbose: + logger.info(f"WebSocketCopilotTarget initialized with conversation_id: {self._conversation_id}") + logger.info(f"Session ID: {self._session_id}") + + @staticmethod + def _dict_to_websocket(data: dict) -> str: + # Produce the smallest possible JSON string, followed by record separator + return json.dumps(data, separators=(",", ":")) + "\x1e" + + @staticmethod + def _parse_message(raw_message: str) -> tuple[int, str, dict]: + """ + Extract actionable content from raw WebSocket frames. + + Args: + raw_message (str): The raw WebSocket message string. + + Returns: + tuple: (message_type, content_text, full_data) + """ + try: + # https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/docs/specs/HubProtocol.md#json-encoding + message = message = raw_message.split("\x1e")[0] # record separator + if not message: + return (-1, "", {}) + + data = json.loads(message) + msg_type = data.get("type", -1) + + if msg_type == 6: # PING + return (6, "", data) + + if msg_type == 2: # LAST_DATA_FRAME + item = data.get("item", {}) + if item: + messages = item.get("messages", []) + if messages: + for msg in reversed(messages): + if msg.get("author") == "bot": + text = msg.get("text", "") + if text: + return (2, text, data) + # TODO: maybe treat this as error? + logger.warning("LAST_DATA_FRAME received but no parseable content found.") + return (2, "", data) + + if msg_type == 1: # NEXT_DATA_FRAME + # Streamed updates are not needed for this target + return (1, "", data) + + return (msg_type, "", data) + + except json.JSONDecodeError as e: + logger.error(f"Failed to decode JSON message: {str(e)}") + return (-1, "", {}) + + def _build_prompt_message(self, prompt: str) -> dict: + return { + "arguments": [ + { + "source": "officeweb", # TODO: support 'teamshub' as well + # TODO: not sure whether to uuid.uuid4() or use a static like it's done in power-pwn + # https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py#L156 + "clientCorrelationId": str(uuid.uuid4()), + "sessionId": self._session_id, + "optionsSets": [ + "enterprise_flux_web", + "enterprise_flux_work", + "enable_request_response_interstitials", + "enterprise_flux_image_v1", + "enterprise_toolbox_with_skdsstore", + "enterprise_toolbox_with_skdsstore_search_message_extensions", + "enable_ME_auth_interstitial", + "skdsstorethirdparty", + "enable_confirmation_interstitial", + "enable_plugin_auth_interstitial", + "enable_response_action_processing", + "enterprise_flux_work_gptv", + "enterprise_flux_work_code_interpreter", + "enable_batch_token_processing", + ], + "options": {}, + "allowedMessageTypes": [ + "Chat", + "Suggestion", + "InternalSearchQuery", + "InternalSearchResult", + "Disengaged", + "InternalLoaderMessage", + "RenderCardRequest", + "AdsQuery", + "SemanticSerp", + "GenerateContentQuery", + "SearchQuery", + "ConfirmationCard", + "AuthError", + "DeveloperLogs", + ], + "sliceIds": [], + # TODO: enable using agents https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py#L192 + "threadLevelGptId": {}, + "conversationId": self._conversation_id, + "traceId": str(uuid.uuid4()).replace("-", ""), # TODO: same case as clientCorrelationId + "isStartOfSession": 0, + "productThreadType": "Office", + "clientInfo": {"clientPlatform": "web"}, + "message": { + "author": "user", + "inputMethod": "Keyboard", + "text": prompt, + "entityAnnotationTypes": ["People", "File", "Event", "Email", "TeamsMessage"], + "requestId": str(uuid.uuid4()).replace("-", ""), + "locationInfo": {"timeZoneOffset": 0, "timeZone": "UTC"}, + "locale": "en-US", + "messageType": "Chat", + "experienceType": "Default", + }, + "plugins": [], # TODO: support enabling some plugins? + } + ], + "invocationId": "0", # TODO: should be dynamic? + "target": "chat", + "type": 4, + } + + async def _connect_and_send(self, prompt: str) -> str: + protocol_msg = {"protocol": "json", "version": 1} + prompt_dict = self._build_prompt_message(prompt) + + inputs = [protocol_msg, prompt_dict] + last_response = "" + + async with websockets.connect(self._websocket_url) as websocket: + for input_msg in inputs: + payload = self._dict_to_websocket(input_msg) + is_user_input = input_msg.get("type") == 4 # USER_PROMPT + + await websocket.send(payload) + + stop_polling = False + while not stop_polling: + response = await websocket.recv() + msg_type, content, data = self._parse_message(response) + + if ( + msg_type in (-1, 2) # UNKNOWN or LAST_DATA_FRAME + or msg_type == 6 + and not is_user_input + ): + stop_polling = True + + if msg_type == 2: # LAST_DATA_FRAME - final response + last_response = content + elif msg_type == -1: # UNKNOWN/NONE + logger.debug("Received unknown or empty message type.") + + return last_response + + def _validate_request(self, *, message: Message) -> None: + n_pieces = len(message.message_pieces) + if n_pieces != 1: + raise ValueError(f"This target only supports a single message piece. Received: {n_pieces} pieces.") + + piece_type = message.message_pieces[0].converted_value_data_type + if piece_type != "text": + raise ValueError(f"This target only supports text prompt input. Received: {piece_type}.") + + @limit_requests_per_minute + async def send_prompt_async(self, *, message: Message) -> list[Message]: + """ + Asynchronously send a message to Microsoft Copilot using WebSocket. + + Args: + message (Message): A message to be sent to the target. + + Returns: + list[Message]: A list containing the response from Copilot. + + Raises: + RuntimeError: If an error occurs during WebSocket communication. + """ + self._validate_request(message=message) + request_piece = message.message_pieces[0] + + try: + prompt_text = request_piece.converted_value + response_text = await self._connect_and_send(prompt_text) + + response_entry = construct_response_from_request( + request=request_piece, response_text_pieces=[response_text] + ) + + return [response_entry] + + except Exception as e: + raise RuntimeError(f"An error occurred during WebSocket communication: {str(e)}") from e diff --git a/websocket_copilot_simple_example.py b/websocket_copilot_simple_example.py new file mode 100644 index 000000000..a1e13831a --- /dev/null +++ b/websocket_copilot_simple_example.py @@ -0,0 +1,31 @@ +""" +# TODO +THIS WILL BE REMOVED after proper unit tests are in place :) +""" + +import asyncio + +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import WebSocketCopilotTarget +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + + +async def main(): + await initialize_pyrit_async(memory_db_type=IN_MEMORY) + target = WebSocketCopilotTarget() + + message_piece = MessagePiece( + role="user", + original_value="say only one random word", + original_value_data_type="text", + converted_value_data_type="text", + ) + message = Message(message_pieces=[message_piece]) + + responses = await target.send_prompt_async(message=message) + for response in responses: + print(f"{response.get_value()}") + + +if __name__ == "__main__": + asyncio.run(main()) From 62fc335afce98e2769bbf0d54bb7e02cb91f33ae Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Tue, 23 Dec 2025 21:07:33 +0100 Subject: [PATCH 02/26] add useful error message for "server rejected WebSocket connection: HTTP 401" --- pyrit/prompt_target/websocket_copilot_target.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index f15994f56..da2b71a67 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -54,7 +54,6 @@ class WebSocketCopilotTarget(PromptTarget): """ # TODO: add more flexible auth, use puppeteer? https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py#L248 - # TODO: add useful message for: "Error during WebSocket communication: server rejected WebSocket connection: HTTP 401" SUPPORTED_DATA_TYPES = {"text"} # TODO: support more types? @@ -278,7 +277,8 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: list[Message]: A list containing the response from Copilot. Raises: - RuntimeError: If an error occurs during WebSocket communication. + websockets.exceptions.InvalidStatus: If the WebSocket connection fails. + RuntimeError: If any other error occurs during WebSocket communication. """ self._validate_request(message=message) request_piece = message.message_pieces[0] @@ -293,5 +293,13 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: return [response_entry] + except websockets.exceptions.InvalidStatus as e: + logger.error( + f"WebSocket connection failed: {str(e)}\n" + "Ensure the WEBSOCKET_URL environment variable is correct and valid." + " For more details about authentication, refer to the class documentation." + ) + raise e + except Exception as e: raise RuntimeError(f"An error occurred during WebSocket communication: {str(e)}") from e From 107b715e920441ed512fc6086a2182dc55a18692 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Tue, 23 Dec 2025 21:38:42 +0100 Subject: [PATCH 03/26] improve error handling and logging --- .../prompt_target/websocket_copilot_target.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index da2b71a67..13669bffb 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -10,6 +10,10 @@ import websockets +from pyrit.exceptions import ( + EmptyResponseException, + pyrit_target_retry, +) from pyrit.models import Message, construct_response_from_request from pyrit.prompt_target import PromptTarget, limit_requests_per_minute @@ -266,6 +270,7 @@ def _validate_request(self, *, message: Message) -> None: raise ValueError(f"This target only supports text prompt input. Received: {piece_type}.") @limit_requests_per_minute + @pyrit_target_retry async def send_prompt_async(self, *, message: Message) -> list[Message]: """ Asynchronously send a message to Microsoft Copilot using WebSocket. @@ -277,16 +282,25 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: list[Message]: A list containing the response from Copilot. Raises: - websockets.exceptions.InvalidStatus: If the WebSocket connection fails. + EmptyResponseException: If the response from Copilot is empty. + InvalidStatus: If the WebSocket handshake fails with an HTTP status error. + WebSocketException: If the WebSocket connection fails. RuntimeError: If any other error occurs during WebSocket communication. """ self._validate_request(message=message) request_piece = message.message_pieces[0] + logger.info(f"Sending the following prompt to WebSocketCopilotTarget: {request_piece}") + try: prompt_text = request_piece.converted_value response_text = await self._connect_and_send(prompt_text) + if not response_text or not response_text.strip(): + logger.error("Empty response received from Copilot.") + raise EmptyResponseException(message="Copilot returned an empty response.") + logger.info(f"Received the following response from WebSocketCopilotTarget: {response_text[:100]}...") + response_entry = construct_response_from_request( request=request_piece, response_text_pieces=[response_text] ) @@ -299,7 +313,9 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: "Ensure the WEBSOCKET_URL environment variable is correct and valid." " For more details about authentication, refer to the class documentation." ) - raise e + raise + except websockets.exceptions.WebSocketException as e: + raise RuntimeError(f"WebSocket communication error: {str(e)}") from e except Exception as e: - raise RuntimeError(f"An error occurred during WebSocket communication: {str(e)}") from e + raise RuntimeError(f"Unexpected error during WebSocket communication: {str(e)}") from e From 41013a22c66c5209e78fa7b62f77300757cf0862 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Tue, 23 Dec 2025 21:51:53 +0100 Subject: [PATCH 04/26] enhance WebSocket URL validation --- .../prompt_target/websocket_copilot_target.py | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 13669bffb..87436004b 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -81,19 +81,30 @@ def __init__( model_name (str): The model name. Defaults to "copilot". Raises: - ValueError: If WebSocket URL is not provided as env variable. + ValueError: If WebSocket URL is not provided, is empty, or has invalid format. + ValueError: If required parameters are missing or empty in the WebSocket URL. """ self._websocket_url = os.getenv("WEBSOCKET_URL") - if not self._websocket_url: + if not self._websocket_url or self._websocket_url.strip() == "": raise ValueError("WebSocket URL must be provided through the WEBSOCKET_URL environment variable") - if not "ConversationId=" in self._websocket_url: - raise ValueError("`ConversationId` parameter not found in URL.") + if not self._websocket_url.startswith(("wss://", "ws://")): + raise ValueError( + "WebSocket URL must start with 'wss://' or 'ws://'. " + f"Received URL starting with: {self._websocket_url[:10]}" + ) + + if "ConversationId=" not in self._websocket_url: + raise ValueError("`ConversationId` parameter not found in WebSocket URL.") self._conversation_id = self._websocket_url.split("ConversationId=")[1].split("&")[0] + if not self._conversation_id: + raise ValueError("`ConversationId` parameter is empty in WebSocket URL.") - if not "X-SessionId=" in self._websocket_url: - raise ValueError("`X-SessionId` parameter not found in URL.") + if "X-SessionId=" not in self._websocket_url: + raise ValueError("`X-SessionId` parameter not found in WebSocket URL.") self._session_id = self._websocket_url.split("X-SessionId=")[1].split("&")[0] + if not self._session_id: + raise ValueError("`X-SessionId` parameter is empty in WebSocket URL.") super().__init__( verbose=verbose, From c8e7e831bd494146076b018ae8d63fd69f374973 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Tue, 23 Dec 2025 22:57:35 +0100 Subject: [PATCH 05/26] small fix --- pyrit/prompt_target/websocket_copilot_target.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 87436004b..a2622e362 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -295,7 +295,6 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: Raises: EmptyResponseException: If the response from Copilot is empty. InvalidStatus: If the WebSocket handshake fails with an HTTP status error. - WebSocketException: If the WebSocket connection fails. RuntimeError: If any other error occurs during WebSocket communication. """ self._validate_request(message=message) @@ -326,7 +325,5 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: ) raise - except websockets.exceptions.WebSocketException as e: - raise RuntimeError(f"WebSocket communication error: {str(e)}") from e except Exception as e: - raise RuntimeError(f"Unexpected error during WebSocket communication: {str(e)}") from e + raise RuntimeError(f"An error occurred during WebSocket communication: {str(e)}") from e From af636b22f9792b1a125cb452963cb03d07df8a1f Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Tue, 23 Dec 2025 23:01:21 +0100 Subject: [PATCH 06/26] fix --- pyrit/prompt_target/websocket_copilot_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index a2622e362..c3b3ebffc 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -135,7 +135,7 @@ def _parse_message(raw_message: str) -> tuple[int, str, dict]: """ try: # https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/docs/specs/HubProtocol.md#json-encoding - message = message = raw_message.split("\x1e")[0] # record separator + message = raw_message.split("\x1e")[0] # record separator if not message: return (-1, "", {}) From c9ebcee1c923f969613bbbfe3077728fe7feaf2a Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Tue, 23 Dec 2025 23:25:09 +0100 Subject: [PATCH 07/26] improve `_parse_message` --- .../prompt_target/websocket_copilot_target.py | 39 ++++++++++--------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index c3b3ebffc..b8f86e076 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -123,7 +123,7 @@ def _dict_to_websocket(data: dict) -> str: return json.dumps(data, separators=(",", ":")) + "\x1e" @staticmethod - def _parse_message(raw_message: str) -> tuple[int, str, dict]: + def _parse_message(raw_message: str) -> tuple[int, str]: """ Extract actionable content from raw WebSocket frames. @@ -131,43 +131,38 @@ def _parse_message(raw_message: str) -> tuple[int, str, dict]: raw_message (str): The raw WebSocket message string. Returns: - tuple: (message_type, content_text, full_data) + tuple[int, str]: A tuple containing the message type and extracted content. """ try: # https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/docs/specs/HubProtocol.md#json-encoding message = raw_message.split("\x1e")[0] # record separator if not message: - return (-1, "", {}) + return (-1, "") data = json.loads(message) msg_type = data.get("type", -1) - if msg_type == 6: # PING - return (6, "", data) + if msg_type in (6, 1): # PING/NEXT_DATA_FRAME + return (msg_type, "") if msg_type == 2: # LAST_DATA_FRAME item = data.get("item", {}) - if item: + if item and isinstance(item, dict): messages = item.get("messages", []) - if messages: + if messages and isinstance(messages, list): for msg in reversed(messages): - if msg.get("author") == "bot": + if isinstance(msg, dict) and msg.get("author") == "bot": text = msg.get("text", "") - if text: - return (2, text, data) - # TODO: maybe treat this as error? + if text and isinstance(text, str): + return (2, text) logger.warning("LAST_DATA_FRAME received but no parseable content found.") - return (2, "", data) + return (2, "") - if msg_type == 1: # NEXT_DATA_FRAME - # Streamed updates are not needed for this target - return (1, "", data) - - return (msg_type, "", data) + return (msg_type, "") except json.JSONDecodeError as e: logger.error(f"Failed to decode JSON message: {str(e)}") - return (-1, "", {}) + return (-1, "") def _build_prompt_message(self, prompt: str) -> dict: return { @@ -255,7 +250,13 @@ async def _connect_and_send(self, prompt: str) -> str: stop_polling = False while not stop_polling: response = await websocket.recv() - msg_type, content, data = self._parse_message(response) + + if response is None: + raise RuntimeError( + "WebSocket connection closed unexpectedly: received None from websocket.recv()" + ) + + msg_type, content = self._parse_message(response) if ( msg_type in (-1, 2) # UNKNOWN or LAST_DATA_FRAME From 99040d0e94238e7dbe61523e7bc23f4549d5f330 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Tue, 23 Dec 2025 23:33:08 +0100 Subject: [PATCH 08/26] useful links --- pyrit/prompt_target/websocket_copilot_target.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index b8f86e076..06d848dbb 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -19,11 +19,9 @@ logger = logging.getLogger(__name__) -""" -Useful links: -https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py -https://labs.zenity.io/p/access-copilot-m365-terminal -""" +# Useful links: +# https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py +# https://labs.zenity.io/p/access-copilot-m365-terminal class CopilotMessageType(Enum): From 1f21ebf87e6c14a845f484e34bd5f97b1cce02f6 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Wed, 24 Dec 2025 23:19:24 +0100 Subject: [PATCH 09/26] add timeouts for responses and connection --- .../prompt_target/websocket_copilot_target.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 06d848dbb..3a51c3009 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import json import logging import os @@ -59,9 +60,8 @@ class WebSocketCopilotTarget(PromptTarget): SUPPORTED_DATA_TYPES = {"text"} # TODO: support more types? - # TODO: implement timeouts and retries - MAX_WAIT_TIME_SECONDS: int = 300 - POLL_INTERVAL_MS: int = 2000 + RESPONSE_TIMEOUT_SECONDS: int = 60 + CONNECTION_TIMEOUT_SECONDS: int = 30 def __init__( self, @@ -238,7 +238,11 @@ async def _connect_and_send(self, prompt: str) -> str: inputs = [protocol_msg, prompt_dict] last_response = "" - async with websockets.connect(self._websocket_url) as websocket: + async with websockets.connect( + self._websocket_url, + open_timeout=self.CONNECTION_TIMEOUT_SECONDS, + close_timeout=self.CONNECTION_TIMEOUT_SECONDS, + ) as websocket: for input_msg in inputs: payload = self._dict_to_websocket(input_msg) is_user_input = input_msg.get("type") == 4 # USER_PROMPT @@ -247,7 +251,15 @@ async def _connect_and_send(self, prompt: str) -> str: stop_polling = False while not stop_polling: - response = await websocket.recv() + try: + response = await asyncio.wait_for( + websocket.recv(), + timeout=self.RECEIVE_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + raise TimeoutError( + f"Timed out waiting for Copilot response after {self.RECEIVE_TIMEOUT_SECONDS} seconds." + ) if response is None: raise RuntimeError( From 1806e7970f00977e241f49a70f50e9b999aa8c1b Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 15:28:10 +0100 Subject: [PATCH 10/26] start with tests --- .../target/test_websocket_copilot_target.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 tests/unit/target/test_websocket_copilot_target.py diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py new file mode 100644 index 000000000..ab45dbb8f --- /dev/null +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +from unittest.mock import patch + +import pytest + +from pyrit.prompt_target import WebSocketCopilotTarget + + +VALID_WEBSOCKET_URL = ( + "wss://substrate.office.com/m365Copilot/Chathub/test_chat_id" + "?ClientRequestId=test_client_request_id" + "&X-SessionId=test_session_id&token=abc123" + "&ConversationId=test_conversation_id" + "&access_token=test_access_token" + # "&variants=feature.test_feature_one,feature.test_feature_two" + # "&agent=web" + # "&scenario=OfficeWebIncludedCopilot" +) + + +@pytest.fixture +def mock_env_websocket_url(): + """Fixture to set the WEBSOCKET_URL environment variable.""" + with patch.dict(os.environ, {"WEBSOCKET_URL": VALID_WEBSOCKET_URL}): + yield + + +@pytest.mark.usefixtures("patch_central_database") +class TestWebSocketCopilotTargetInit: + def test_init_with_valid_wss_url(self, mock_env_websocket_url): + target = WebSocketCopilotTarget() + + assert target._websocket_url == VALID_WEBSOCKET_URL + assert target._conversation_id == "test_conversation_id" + assert target._session_id == "test_session_id" + assert target._model_name == "copilot" + + def test_init_with_missing_or_invalid_wss_url(self): + for env_vars in [{}, {"WEBSOCKET_URL": ""}, {"WEBSOCKET_URL": " "}]: + with patch.dict(os.environ, env_vars, clear=True): + with pytest.raises(ValueError, match="WebSocket URL must be provided"): + WebSocketCopilotTarget() + + for invalid_url in ["invalid_websocket_url", "ws://example.com", "https://example.com"]: + with patch.dict(os.environ, {"WEBSOCKET_URL": invalid_url}, clear=True): + with pytest.raises(ValueError, match="WebSocket URL must start with 'wss://'"): + WebSocketCopilotTarget() + + def test_init_with_missing_or_empty_required_params(self): + urls = [ + ("wss://example.com/?X-SessionId=session123", "`ConversationId` parameter not found"), + ("wss://example.com/?ConversationId=conv123", "`X-SessionId` parameter not found"), + ("wss://example.com/?ConversationId=&X-SessionId=session123", "`ConversationId` parameter is empty"), + ("wss://example.com/?ConversationId=conv123&X-SessionId=", "`X-SessionId` parameter is empty"), + ] + + for url, error_msg in urls: + with patch.dict(os.environ, {"WEBSOCKET_URL": url}, clear=True): + with pytest.raises(ValueError, match=error_msg): + WebSocketCopilotTarget() + + def test_init_sets_endpoint_correctly(self, mock_env_websocket_url): + target = WebSocketCopilotTarget() + assert target._endpoint == "wss://substrate.office.com/m365Copilot/Chathub/test_chat_id" From 3962bbbab536d45accd9c2a413215394be05d546 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 15:28:39 +0100 Subject: [PATCH 11/26] require `wss://` only --- pyrit/prompt_target/websocket_copilot_target.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 3a51c3009..ea99198a9 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -86,11 +86,8 @@ def __init__( if not self._websocket_url or self._websocket_url.strip() == "": raise ValueError("WebSocket URL must be provided through the WEBSOCKET_URL environment variable") - if not self._websocket_url.startswith(("wss://", "ws://")): - raise ValueError( - "WebSocket URL must start with 'wss://' or 'ws://'. " - f"Received URL starting with: {self._websocket_url[:10]}" - ) + if not self._websocket_url.startswith("wss://"): + raise ValueError(f"WebSocket URL must start with 'wss://'. Received: {self._websocket_url[:10]}") if "ConversationId=" not in self._websocket_url: raise ValueError("`ConversationId` parameter not found in WebSocket URL.") From 7588e8d1917dc5947c4f2bb4cff085cf34b2bd1d Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 15:41:24 +0100 Subject: [PATCH 12/26] add configurable response timeout --- pyrit/prompt_target/websocket_copilot_target.py | 10 ++++++++-- tests/unit/target/test_websocket_copilot_target.py | 9 +++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index ea99198a9..decd8c453 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -69,6 +69,7 @@ def __init__( verbose: bool = False, max_requests_per_minute: Optional[int] = None, model_name: str = "copilot", + response_timeout_seconds: int = RESPONSE_TIMEOUT_SECONDS, ) -> None: """ Initialize the WebSocketCopilotTarget. @@ -77,6 +78,7 @@ def __init__( verbose (bool): Enable verbose logging. Defaults to False. max_requests_per_minute (int, Optional): 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. Raises: ValueError: If WebSocket URL is not provided, is empty, or has invalid format. @@ -108,6 +110,10 @@ def __init__( model_name=model_name, ) + if response_timeout_seconds <= 0: + raise ValueError("response_timeout_seconds must be a positive integer.") + self._response_timeout_seconds = response_timeout_seconds + if self._verbose: logger.info(f"WebSocketCopilotTarget initialized with conversation_id: {self._conversation_id}") logger.info(f"Session ID: {self._session_id}") @@ -251,11 +257,11 @@ async def _connect_and_send(self, prompt: str) -> str: try: response = await asyncio.wait_for( websocket.recv(), - timeout=self.RECEIVE_TIMEOUT_SECONDS, + timeout=self._response_timeout_seconds, ) except asyncio.TimeoutError: raise TimeoutError( - f"Timed out waiting for Copilot response after {self.RECEIVE_TIMEOUT_SECONDS} seconds." + f"Timed out waiting for Copilot response after {self._response_timeout_seconds} seconds." ) if response is None: diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index ab45dbb8f..37fe2f345 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -65,3 +65,12 @@ def test_init_with_missing_or_empty_required_params(self): def test_init_sets_endpoint_correctly(self, mock_env_websocket_url): target = WebSocketCopilotTarget() assert target._endpoint == "wss://substrate.office.com/m365Copilot/Chathub/test_chat_id" + + def test_init_with_custom_response_timeout(self, mock_env_websocket_url): + target = WebSocketCopilotTarget(response_timeout_seconds=120) + assert target._response_timeout_seconds == 120 + + for invalid_timeout in [0, -10]: + with pytest.raises(ValueError, match="response_timeout_seconds must be a positive integer."): + WebSocketCopilotTarget(response_timeout_seconds=invalid_timeout) + From b98f6754d7e4d68091fddbacb31982f35284e0a1 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 15:43:27 +0100 Subject: [PATCH 13/26] fix --- tests/unit/target/test_websocket_copilot_target.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index 37fe2f345..deae55ca2 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -30,7 +30,7 @@ def mock_env_websocket_url(): @pytest.mark.usefixtures("patch_central_database") class TestWebSocketCopilotTargetInit: - def test_init_with_valid_wss_url(self, mock_env_websocket_url): + def test_init_with_valid_wss_url(self): target = WebSocketCopilotTarget() assert target._websocket_url == VALID_WEBSOCKET_URL @@ -62,11 +62,11 @@ def test_init_with_missing_or_empty_required_params(self): with pytest.raises(ValueError, match=error_msg): WebSocketCopilotTarget() - def test_init_sets_endpoint_correctly(self, mock_env_websocket_url): + def test_init_sets_endpoint_correctly(self): target = WebSocketCopilotTarget() assert target._endpoint == "wss://substrate.office.com/m365Copilot/Chathub/test_chat_id" - def test_init_with_custom_response_timeout(self, mock_env_websocket_url): + def test_init_with_custom_response_timeout(self): target = WebSocketCopilotTarget(response_timeout_seconds=120) assert target._response_timeout_seconds == 120 From 0dab7bb0eefe1834bc7b28cf07b2e9bbffa9b82f Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 16:02:50 +0100 Subject: [PATCH 14/26] replace Enum with IntEnum and actually use it --- .../prompt_target/websocket_copilot_target.py | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index decd8c453..9af6f7fd1 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -6,7 +6,7 @@ import logging import os import uuid -from enum import Enum +from enum import IntEnum from typing import Optional import websockets @@ -25,7 +25,7 @@ # https://labs.zenity.io/p/access-copilot-m365-terminal -class CopilotMessageType(Enum): +class CopilotMessageType(IntEnum): """Enumeration for Copilot WebSocket message types.""" UNKNOWN = -1 @@ -124,7 +124,7 @@ def _dict_to_websocket(data: dict) -> str: return json.dumps(data, separators=(",", ":")) + "\x1e" @staticmethod - def _parse_message(raw_message: str) -> tuple[int, str]: + def _parse_message(raw_message: str) -> tuple[CopilotMessageType, str]: """ Extract actionable content from raw WebSocket frames. @@ -132,21 +132,21 @@ def _parse_message(raw_message: str) -> tuple[int, str]: raw_message (str): The raw WebSocket message string. Returns: - tuple[int, str]: A tuple containing the message type and extracted content. + tuple[CopilotMessageType, str]: A tuple containing the message type and extracted content. """ try: # https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/docs/specs/HubProtocol.md#json-encoding message = raw_message.split("\x1e")[0] # record separator if not message: - return (-1, "") + return (CopilotMessageType.UNKNOWN, "") data = json.loads(message) - msg_type = data.get("type", -1) + msg_type = CopilotMessageType(data.get("type", -1)) - if msg_type in (6, 1): # PING/NEXT_DATA_FRAME + if msg_type in (CopilotMessageType.PING, CopilotMessageType.NEXT_DATA_FRAME): return (msg_type, "") - if msg_type == 2: # LAST_DATA_FRAME + if msg_type == CopilotMessageType.LAST_DATA_FRAME: item = data.get("item", {}) if item and isinstance(item, dict): messages = item.get("messages", []) @@ -155,15 +155,15 @@ def _parse_message(raw_message: str) -> tuple[int, str]: if isinstance(msg, dict) and msg.get("author") == "bot": text = msg.get("text", "") if text and isinstance(text, str): - return (2, text) + return (CopilotMessageType.LAST_DATA_FRAME, text) logger.warning("LAST_DATA_FRAME received but no parseable content found.") - return (2, "") + return (CopilotMessageType.LAST_DATA_FRAME, "") return (msg_type, "") except json.JSONDecodeError as e: logger.error(f"Failed to decode JSON message: {str(e)}") - return (-1, "") + return (CopilotMessageType.UNKNOWN, "") def _build_prompt_message(self, prompt: str) -> dict: return { @@ -231,7 +231,7 @@ def _build_prompt_message(self, prompt: str) -> dict: ], "invocationId": "0", # TODO: should be dynamic? "target": "chat", - "type": 4, + "type": CopilotMessageType.USER_PROMPT, } async def _connect_and_send(self, prompt: str) -> str: @@ -248,7 +248,7 @@ async def _connect_and_send(self, prompt: str) -> str: ) as websocket: for input_msg in inputs: payload = self._dict_to_websocket(input_msg) - is_user_input = input_msg.get("type") == 4 # USER_PROMPT + is_user_input = input_msg.get("type") == CopilotMessageType.USER_PROMPT await websocket.send(payload) @@ -272,15 +272,15 @@ async def _connect_and_send(self, prompt: str) -> str: msg_type, content = self._parse_message(response) if ( - msg_type in (-1, 2) # UNKNOWN or LAST_DATA_FRAME - or msg_type == 6 + msg_type in (CopilotMessageType.UNKNOWN, CopilotMessageType.LAST_DATA_FRAME) + or msg_type == CopilotMessageType.PING and not is_user_input ): stop_polling = True - if msg_type == 2: # LAST_DATA_FRAME - final response + if msg_type == CopilotMessageType.LAST_DATA_FRAME: last_response = content - elif msg_type == -1: # UNKNOWN/NONE + elif msg_type == CopilotMessageType.UNKNOWN: logger.debug("Received unknown or empty message type.") return last_response From c2df619c5727593068012ebb5eeea46ffeba42a4 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 16:18:35 +0100 Subject: [PATCH 15/26] test_dict_to_websocket_static_method --- .../target/test_websocket_copilot_target.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index deae55ca2..663ce5145 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -21,13 +21,6 @@ ) -@pytest.fixture -def mock_env_websocket_url(): - """Fixture to set the WEBSOCKET_URL environment variable.""" - with patch.dict(os.environ, {"WEBSOCKET_URL": VALID_WEBSOCKET_URL}): - yield - - @pytest.mark.usefixtures("patch_central_database") class TestWebSocketCopilotTargetInit: def test_init_with_valid_wss_url(self): @@ -74,3 +67,16 @@ def test_init_with_custom_response_timeout(self): with pytest.raises(ValueError, match="response_timeout_seconds must be a positive integer."): WebSocketCopilotTarget(response_timeout_seconds=invalid_timeout) + +@pytest.mark.parametrize( + "data,expected", + [ + ({"key": "value"}, '{"key":"value"}\x1e'), + ({"protocol": "json", "version": 1}, '{"protocol":"json","version":1}\x1e'), + ({"outer": {"inner": "value"}}, '{"outer":{"inner":"value"}}\x1e'), + ({"items": [1, 2, 3]}, '{"items":[1,2,3]}\x1e'), + ], +) +def test_dict_to_websocket_static_method(data, expected): + result = WebSocketCopilotTarget._dict_to_websocket(data) + assert result == expected From 18fd2387d0cb9c06df15a2bd0f6a600cc3bb675e Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 16:20:56 +0100 Subject: [PATCH 16/26] fix --- tests/unit/target/test_websocket_copilot_target.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index 663ce5145..aadbc1ed5 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -21,9 +21,15 @@ ) +@pytest.fixture +def mock_env_websocket_url(): + with patch.dict(os.environ, {"WEBSOCKET_URL": VALID_WEBSOCKET_URL}): + yield + + @pytest.mark.usefixtures("patch_central_database") class TestWebSocketCopilotTargetInit: - def test_init_with_valid_wss_url(self): + def test_init_with_valid_wss_url(self, mock_env_websocket_url): target = WebSocketCopilotTarget() assert target._websocket_url == VALID_WEBSOCKET_URL @@ -55,11 +61,11 @@ def test_init_with_missing_or_empty_required_params(self): with pytest.raises(ValueError, match=error_msg): WebSocketCopilotTarget() - def test_init_sets_endpoint_correctly(self): + def test_init_sets_endpoint_correctly(self, mock_env_websocket_url): target = WebSocketCopilotTarget() assert target._endpoint == "wss://substrate.office.com/m365Copilot/Chathub/test_chat_id" - def test_init_with_custom_response_timeout(self): + def test_init_with_custom_response_timeout(self, mock_env_websocket_url): target = WebSocketCopilotTarget(response_timeout_seconds=120) assert target._response_timeout_seconds == 120 From 73b07d00272adaa089acf888fb57bb954ffb487d Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 18:04:05 +0100 Subject: [PATCH 17/26] Refactor WebSocket message parser to handle multiple frames per message - Rename _parse_message() to _parse_raw_message() - Split on record separator (\x1e) and processes all frames, not just first - Add FINAL_DATA_FRAME (type 3) enum value for completion signals - Extract bot message parsing logic into lambda for reusability - Fixes stop condition to handle FINAL_DATA_FRAME and remove flawed "is_user_input and PING" logic --- .../prompt_target/websocket_copilot_target.py | 111 ++++++++++-------- 1 file changed, 64 insertions(+), 47 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 9af6f7fd1..c55ea051a 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -31,6 +31,7 @@ class CopilotMessageType(IntEnum): UNKNOWN = -1 NEXT_DATA_FRAME = 1 # streaming Copilot responses LAST_DATA_FRAME = 2 # the last data frame with final content + FINAL_DATA_FRAME = 3 # the final data frame indicating completion USER_PROMPT = 4 PING = 6 @@ -124,46 +125,63 @@ def _dict_to_websocket(data: dict) -> str: return json.dumps(data, separators=(",", ":")) + "\x1e" @staticmethod - def _parse_message(raw_message: str) -> tuple[CopilotMessageType, str]: + def _parse_raw_message(message: str) -> list[tuple[CopilotMessageType, str]]: """ - Extract actionable content from raw WebSocket frames. + Extract actionable content from a raw WebSocket message. + Returns more than one JSON message if multiple are found. Args: - raw_message (str): The raw WebSocket message string. + message (str): The raw WebSocket message string. Returns: - tuple[CopilotMessageType, str]: A tuple containing the message type and extracted content. + list[tuple[CopilotMessageType, str]]: A list of tuples where each tuple contains + message type and extracted content. """ - try: - # https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/docs/specs/HubProtocol.md#json-encoding - message = raw_message.split("\x1e")[0] # record separator - if not message: - return (CopilotMessageType.UNKNOWN, "") - - data = json.loads(message) - msg_type = CopilotMessageType(data.get("type", -1)) - - if msg_type in (CopilotMessageType.PING, CopilotMessageType.NEXT_DATA_FRAME): - return (msg_type, "") - - if msg_type == CopilotMessageType.LAST_DATA_FRAME: - item = data.get("item", {}) - if item and isinstance(item, dict): - messages = item.get("messages", []) - if messages and isinstance(messages, list): - for msg in reversed(messages): - if isinstance(msg, dict) and msg.get("author") == "bot": - text = msg.get("text", "") - if text and isinstance(text, str): - return (CopilotMessageType.LAST_DATA_FRAME, text) - logger.warning("LAST_DATA_FRAME received but no parseable content found.") - return (CopilotMessageType.LAST_DATA_FRAME, "") - - return (msg_type, "") - - except json.JSONDecodeError as e: - logger.error(f"Failed to decode JSON message: {str(e)}") - return (CopilotMessageType.UNKNOWN, "") + # Find the last chat message with text content + extract_bot_message = lambda data: next( + ( + msg.get("text", "") + for msg in reversed(data.get("item", {}).get("messages", [])) + if isinstance(msg, dict) and msg.get("author") == "bot" and msg.get("text") + ), + "", + ) + + results: list[tuple[CopilotMessageType, str]] = [] + + # https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/docs/specs/HubProtocol.md#json-encoding + messages = message.split("\x1e") # record separator + + for message in messages: + if not message or not message.strip(): + continue + + try: + data = json.loads(message) + msg_type = CopilotMessageType(data.get("type", -1)) + + if msg_type in ( + CopilotMessageType.PING, + CopilotMessageType.NEXT_DATA_FRAME, + CopilotMessageType.FINAL_DATA_FRAME, + ): + results.append((msg_type, "")) + continue + + if msg_type == CopilotMessageType.LAST_DATA_FRAME: + bot_text = extract_bot_message(data) + if not bot_text: + logger.warning("LAST_DATA_FRAME received but no parseable content found.") + results.append((CopilotMessageType.LAST_DATA_FRAME, bot_text)) + continue + + results.append((msg_type, "")) + + except json.JSONDecodeError as e: + logger.error(f"Failed to decode JSON message: {str(e)}") + results.append((CopilotMessageType.UNKNOWN, "")) + + return results if results else [(CopilotMessageType.UNKNOWN, "")] def _build_prompt_message(self, prompt: str) -> dict: return { @@ -248,8 +266,6 @@ async def _connect_and_send(self, prompt: str) -> str: ) as websocket: for input_msg in inputs: payload = self._dict_to_websocket(input_msg) - is_user_input = input_msg.get("type") == CopilotMessageType.USER_PROMPT - await websocket.send(payload) stop_polling = False @@ -269,19 +285,20 @@ async def _connect_and_send(self, prompt: str) -> str: "WebSocket connection closed unexpectedly: received None from websocket.recv()" ) - msg_type, content = self._parse_message(response) + parsed_messages = self._parse_raw_message(response) - if ( - msg_type in (CopilotMessageType.UNKNOWN, CopilotMessageType.LAST_DATA_FRAME) - or msg_type == CopilotMessageType.PING - and not is_user_input - ): - stop_polling = True + for msg_type, content in parsed_messages: + if msg_type in ( + CopilotMessageType.UNKNOWN, + CopilotMessageType.LAST_DATA_FRAME, + CopilotMessageType.FINAL_DATA_FRAME, + ): + stop_polling = True - if msg_type == CopilotMessageType.LAST_DATA_FRAME: - last_response = content - elif msg_type == CopilotMessageType.UNKNOWN: - logger.debug("Received unknown or empty message type.") + if msg_type == CopilotMessageType.LAST_DATA_FRAME: + last_response = content + elif msg_type == CopilotMessageType.UNKNOWN: + logger.debug("Received unknown or empty message type.") return last_response From 9a8a878ebc45c9fc7c7f251a449ca401f6246b1c Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 18:10:05 +0100 Subject: [PATCH 18/26] rename message types in the enum --- .../prompt_target/websocket_copilot_target.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index c55ea051a..5956154d5 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -29,9 +29,9 @@ class CopilotMessageType(IntEnum): """Enumeration for Copilot WebSocket message types.""" UNKNOWN = -1 - NEXT_DATA_FRAME = 1 # streaming Copilot responses - LAST_DATA_FRAME = 2 # the last data frame with final content - FINAL_DATA_FRAME = 3 # the final data frame indicating completion + PARTIAL_RESPONSE = 1 + FINAL_CONTENT = 2 + STREAM_END = 3 USER_PROMPT = 4 PING = 6 @@ -162,17 +162,17 @@ def _parse_raw_message(message: str) -> list[tuple[CopilotMessageType, str]]: if msg_type in ( CopilotMessageType.PING, - CopilotMessageType.NEXT_DATA_FRAME, - CopilotMessageType.FINAL_DATA_FRAME, + CopilotMessageType.PARTIAL_RESPONSE, + CopilotMessageType.STREAM_END, ): results.append((msg_type, "")) continue - if msg_type == CopilotMessageType.LAST_DATA_FRAME: + if msg_type == CopilotMessageType.FINAL_CONTENT: bot_text = extract_bot_message(data) if not bot_text: - logger.warning("LAST_DATA_FRAME received but no parseable content found.") - results.append((CopilotMessageType.LAST_DATA_FRAME, bot_text)) + logger.warning("FINAL_CONTENT received but no parseable content found.") + results.append((CopilotMessageType.FINAL_CONTENT, bot_text)) continue results.append((msg_type, "")) @@ -290,12 +290,12 @@ async def _connect_and_send(self, prompt: str) -> str: for msg_type, content in parsed_messages: if msg_type in ( CopilotMessageType.UNKNOWN, - CopilotMessageType.LAST_DATA_FRAME, - CopilotMessageType.FINAL_DATA_FRAME, + CopilotMessageType.FINAL_CONTENT, + CopilotMessageType.STREAM_END, ): stop_polling = True - if msg_type == CopilotMessageType.LAST_DATA_FRAME: + if msg_type == CopilotMessageType.FINAL_CONTENT: last_response = content elif msg_type == CopilotMessageType.UNKNOWN: logger.debug("Received unknown or empty message type.") From 4d3c15dd1c575a5f48d272847b2d9b1106c4c7bc Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 18:26:34 +0100 Subject: [PATCH 19/26] add raw WebSocket messages for testing --- tests/unit/target/test_websocket_copilot_target.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index aadbc1ed5..a1588b413 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -86,3 +86,14 @@ def test_init_with_custom_response_timeout(self, mock_env_websocket_url): def test_dict_to_websocket_static_method(data, expected): result = WebSocketCopilotTarget._dict_to_websocket(data) assert result == expected + + +RAW_WEBSOCKET_MESSAGES = [ + "{}\x1e", + '{"type":6}\x1e', + '{"type":1,"target":"update","arguments":[{"messages":[{"text":"Apple","author":"bot","responseIdentifier":"Default"}]}]}\x1e', + '{"type":3,"invocationId":"0"}\x1e', + '{"type":2,"invocationId":"0","item":{"messages":[{"text":"Name a fruit","author":"user"},{"text":"Apple. 🍎 \n\nWould you like me to list more fruits or give you some interesting facts about apples?","turnState":"Completed","author":"bot","turnCount":1}],"firstNewMessageIndex":1,"conversationId":"conversationId","requestId":"requestId","result":{"value":"Success","message":"Apple. 🍎 \n\nWould you like me to list more fruits or give you some interesting facts about apples?","serviceVersion":"1.0.03273.12483"}}}\x1e', +] + +# TODO: add tests for _parse_raw_message From b095d742476ae416189dc177815a786a42543e91 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 19:43:20 +0100 Subject: [PATCH 20/26] remove emojis --- tests/unit/target/test_websocket_copilot_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index a1588b413..50b6681fc 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -93,7 +93,7 @@ def test_dict_to_websocket_static_method(data, expected): '{"type":6}\x1e', '{"type":1,"target":"update","arguments":[{"messages":[{"text":"Apple","author":"bot","responseIdentifier":"Default"}]}]}\x1e', '{"type":3,"invocationId":"0"}\x1e', - '{"type":2,"invocationId":"0","item":{"messages":[{"text":"Name a fruit","author":"user"},{"text":"Apple. 🍎 \n\nWould you like me to list more fruits or give you some interesting facts about apples?","turnState":"Completed","author":"bot","turnCount":1}],"firstNewMessageIndex":1,"conversationId":"conversationId","requestId":"requestId","result":{"value":"Success","message":"Apple. 🍎 \n\nWould you like me to list more fruits or give you some interesting facts about apples?","serviceVersion":"1.0.03273.12483"}}}\x1e', + '{"type":2,"invocationId":"0","item":{"messages":[{"text":"Name a fruit","author":"user"},{"text":"Apple. \n\nWould you like me to list more fruits or give you some interesting facts about apples?","turnState":"Completed","author":"bot","turnCount":1}],"firstNewMessageIndex":1,"conversationId":"conversationId","requestId":"requestId","result":{"value":"Success","message":"Apple. \n\nWould you like me to list more fruits or give you some interesting facts about apples?","serviceVersion":"1.0.03273.12483"}}}\x1e', ] # TODO: add tests for _parse_raw_message From 38e686827a05bb9fbb20ad03b09a9dd13da557ec Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 20:12:56 +0100 Subject: [PATCH 21/26] simpler way to get the final result --- pyrit/prompt_target/websocket_copilot_target.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 5956154d5..47631f694 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -137,16 +137,6 @@ def _parse_raw_message(message: str) -> list[tuple[CopilotMessageType, str]]: list[tuple[CopilotMessageType, str]]: A list of tuples where each tuple contains message type and extracted content. """ - # Find the last chat message with text content - extract_bot_message = lambda data: next( - ( - msg.get("text", "") - for msg in reversed(data.get("item", {}).get("messages", [])) - if isinstance(msg, dict) and msg.get("author") == "bot" and msg.get("text") - ), - "", - ) - results: list[tuple[CopilotMessageType, str]] = [] # https://github.com/dotnet/aspnetcore/blob/main/src/SignalR/docs/specs/HubProtocol.md#json-encoding @@ -169,8 +159,9 @@ def _parse_raw_message(message: str) -> list[tuple[CopilotMessageType, str]]: continue if msg_type == CopilotMessageType.FINAL_CONTENT: - bot_text = extract_bot_message(data) + bot_text = data.get("item", {}).get("result", {}).get("message", "") if not bot_text: + # In this case, EmptyResponseException will be raised anyway logger.warning("FINAL_CONTENT received but no parseable content found.") results.append((CopilotMessageType.FINAL_CONTENT, bot_text)) continue From 2430dbecba63b95a34c79efcf161240cb894b4b7 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 20:27:01 +0100 Subject: [PATCH 22/26] log full raw message when no parseable content found --- pyrit/prompt_target/websocket_copilot_target.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 47631f694..fae6fef35 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -163,6 +163,7 @@ def _parse_raw_message(message: str) -> list[tuple[CopilotMessageType, str]]: if not bot_text: # In this case, EmptyResponseException will be raised anyway logger.warning("FINAL_CONTENT received but no parseable content found.") + logger.debug(f"Full raw message: {message}") results.append((CopilotMessageType.FINAL_CONTENT, bot_text)) continue From 5b2c54a47bc8e62e2a1c997ba579a9bd5771a0e9 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 21:50:01 +0100 Subject: [PATCH 23/26] _value2member_map_ --- pyrit/prompt_target/websocket_copilot_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index fae6fef35..58a380ecd 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -148,7 +148,7 @@ def _parse_raw_message(message: str) -> list[tuple[CopilotMessageType, str]]: try: data = json.loads(message) - msg_type = CopilotMessageType(data.get("type", -1)) + msg_type = CopilotMessageType._value2member_map_.get(data.get("type", -1), CopilotMessageType.UNKNOWN) if msg_type in ( CopilotMessageType.PING, From 4a7a7b8bd922f612cb26bf3c2d2ded97c9f33ad9 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 21:50:55 +0100 Subject: [PATCH 24/26] TestParseRawMessage --- .../target/test_websocket_copilot_target.py | 71 ++++++++++++++++--- 1 file changed, 62 insertions(+), 9 deletions(-) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index 50b6681fc..661cede37 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -88,12 +88,65 @@ def test_dict_to_websocket_static_method(data, expected): assert result == expected -RAW_WEBSOCKET_MESSAGES = [ - "{}\x1e", - '{"type":6}\x1e', - '{"type":1,"target":"update","arguments":[{"messages":[{"text":"Apple","author":"bot","responseIdentifier":"Default"}]}]}\x1e', - '{"type":3,"invocationId":"0"}\x1e', - '{"type":2,"invocationId":"0","item":{"messages":[{"text":"Name a fruit","author":"user"},{"text":"Apple. \n\nWould you like me to list more fruits or give you some interesting facts about apples?","turnState":"Completed","author":"bot","turnCount":1}],"firstNewMessageIndex":1,"conversationId":"conversationId","requestId":"requestId","result":{"value":"Success","message":"Apple. \n\nWould you like me to list more fruits or give you some interesting facts about apples?","serviceVersion":"1.0.03273.12483"}}}\x1e', -] - -# TODO: add tests for _parse_raw_message +class TestParseRawMessage: + from pyrit.prompt_target.websocket_copilot_target import CopilotMessageType + + @pytest.mark.parametrize( + "message,expected_types,expected_content", + [ + ("", [CopilotMessageType.UNKNOWN], [""]), + (" \n\t ", [CopilotMessageType.UNKNOWN], [""]), + ("{}\x1e", [CopilotMessageType.UNKNOWN], [""]), + ('{"type":6}\x1e', [CopilotMessageType.PING], [""]), + ( + '{"type":1,"target":"update","arguments":[{"messages":[{"text":"Partial","author":"bot"}]}]}\x1e', + [CopilotMessageType.PARTIAL_RESPONSE], + [""], + ), + ( + '{"type":2,"item":{"result":{"message":"Final."}}}\x1e{"type":3,"invocationId":"0"}\x1e', + [CopilotMessageType.FINAL_CONTENT, CopilotMessageType.STREAM_END], + [ + "Final.", + "", + ], + ), + ], + ) + def test_parse_raw_message_with_valid_data(self, message, expected_types, expected_content): + result = WebSocketCopilotTarget._parse_raw_message(message) + + assert len(result) == len(expected_types) + for i, expected_type in enumerate(expected_types): + assert result[i][0] == expected_type + assert result[i][1] == expected_content[i] + + def test_parse_final_message_without_content(self): + from pyrit.prompt_target.websocket_copilot_target import CopilotMessageType + + with patch("pyrit.prompt_target.websocket_copilot_target.logger") as mock_logger: + message = '{"type":2,"invocationId":"0"}\x1e' + result = WebSocketCopilotTarget._parse_raw_message(message) + + assert len(result) == 1 + assert result[0][0] == CopilotMessageType.FINAL_CONTENT + assert result[0][1] == "" + + mock_logger.warning.assert_called_with("FINAL_CONTENT received but no parseable content found.") + mock_logger.debug.assert_called_with(f"Full raw message: {message[:-1]}") + + @pytest.mark.parametrize( + "message", + [ + '{"type":99,"data":"unknown"}\x1e', + '{"data":"no type field"}\x1e', + '{"invalid json structure\x1e', + ], + ) + def test_parse_unknown_or_invalid_messages(self, message): + from pyrit.prompt_target.websocket_copilot_target import CopilotMessageType + + result = WebSocketCopilotTarget._parse_raw_message(message) + assert len(result) == 1 + assert result[0][0] == CopilotMessageType.UNKNOWN + assert result[0][1] == "" From acb0a6da8f6c162f5887819316542c356c70e167 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 23:30:56 +0100 Subject: [PATCH 25/26] test fix --- tests/unit/target/test_websocket_copilot_target.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/target/test_websocket_copilot_target.py b/tests/unit/target/test_websocket_copilot_target.py index 661cede37..0b4e713e1 100644 --- a/tests/unit/target/test_websocket_copilot_target.py +++ b/tests/unit/target/test_websocket_copilot_target.py @@ -10,7 +10,7 @@ VALID_WEBSOCKET_URL = ( - "wss://substrate.office.com/m365Copilot/Chathub/test_chat_id" + "wss://substrate.office.com/m365Copilot/Chathub/test_object_id@test_tenant_id" "?ClientRequestId=test_client_request_id" "&X-SessionId=test_session_id&token=abc123" "&ConversationId=test_conversation_id" @@ -63,7 +63,7 @@ def test_init_with_missing_or_empty_required_params(self): def test_init_sets_endpoint_correctly(self, mock_env_websocket_url): target = WebSocketCopilotTarget() - assert target._endpoint == "wss://substrate.office.com/m365Copilot/Chathub/test_chat_id" + assert target._endpoint == "wss://substrate.office.com/m365Copilot/Chathub/test_object_id@test_tenant_id" def test_init_with_custom_response_timeout(self, mock_env_websocket_url): target = WebSocketCopilotTarget(response_timeout_seconds=120) From 276290f7dc309731fb1cba38c14132026833ee51 Mon Sep 17 00:00:00 2001 From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com> Date: Thu, 25 Dec 2025 23:32:07 +0100 Subject: [PATCH 26/26] TODO: use msal for auth --- .../prompt_target/websocket_copilot_target.py | 102 ++++++++---------- 1 file changed, 46 insertions(+), 56 deletions(-) diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 58a380ecd..7f6137535 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -57,8 +57,6 @@ class WebSocketCopilotTarget(PromptTarget): Only works with licensed Microsoft 365 Copilot. The free Copilot version is not compatible. """ - # TODO: add more flexible auth, use puppeteer? https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py#L248 - SUPPORTED_DATA_TYPES = {"text"} # TODO: support more types? RESPONSE_TIMEOUT_SECONDS: int = 60 @@ -179,67 +177,59 @@ def _build_prompt_message(self, prompt: str) -> dict: return { "arguments": [ { - "source": "officeweb", # TODO: support 'teamshub' as well - # TODO: not sure whether to uuid.uuid4() or use a static like it's done in power-pwn - # https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py#L156 - "clientCorrelationId": str(uuid.uuid4()), - "sessionId": self._session_id, - "optionsSets": [ - "enterprise_flux_web", - "enterprise_flux_work", - "enable_request_response_interstitials", - "enterprise_flux_image_v1", - "enterprise_toolbox_with_skdsstore", - "enterprise_toolbox_with_skdsstore_search_message_extensions", - "enable_ME_auth_interstitial", - "skdsstorethirdparty", - "enable_confirmation_interstitial", - "enable_plugin_auth_interstitial", - "enable_response_action_processing", - "enterprise_flux_work_gptv", - "enterprise_flux_work_code_interpreter", - "enable_batch_token_processing", - ], - "options": {}, - "allowedMessageTypes": [ - "Chat", - "Suggestion", - "InternalSearchQuery", - "InternalSearchResult", - "Disengaged", - "InternalLoaderMessage", - "RenderCardRequest", - "AdsQuery", - "SemanticSerp", - "GenerateContentQuery", - "SearchQuery", - "ConfirmationCard", - "AuthError", - "DeveloperLogs", - ], - "sliceIds": [], - # TODO: enable using agents https://github.com/mbrg/power-pwn/blob/main/src/powerpwn/copilot/copilot_connector/copilot_connector.py#L192 - "threadLevelGptId": {}, - "conversationId": self._conversation_id, - "traceId": str(uuid.uuid4()).replace("-", ""), # TODO: same case as clientCorrelationId - "isStartOfSession": 0, - "productThreadType": "Office", - "clientInfo": {"clientPlatform": "web"}, + # TODO: use msal for auth, then set these fields properly, as with current approach they are not really needed + # "source": "officeweb", + # "clientCorrelationId": str(uuid.uuid4()), + # "sessionId": self._session_id, + # "optionsSets": [ + # "enterprise_flux_web", + # "enterprise_flux_work", + # "enable_request_response_interstitials", + # "enterprise_flux_image_v1", + # "enterprise_toolbox_with_skdsstore", + # "enterprise_toolbox_with_skdsstore_search_message_extensions", + # "enable_ME_auth_interstitial", + # "skdsstorethirdparty", + # "enable_confirmation_interstitial", + # "enable_plugin_auth_interstitial", + # "enable_response_action_processing", + # "enterprise_flux_work_gptv", + # "enterprise_flux_work_code_interpreter", + # "enable_batch_token_processing", + # ], + # "options": {}, + # "allowedMessageTypes": [ + # "Chat", + # "Suggestion", + # "InternalSearchQuery", + # "InternalSearchResult", + # "Disengaged", + # "InternalLoaderMessage", + # "RenderCardRequest", + # "AdsQuery", + # "SemanticSerp", + # "GenerateContentQuery", + # "SearchQuery", + # "ConfirmationCard", + # "AuthError", + # "DeveloperLogs", + # ], + # "sliceIds": [], + # "threadLevelGptId": {}, + # "conversationId": self._conversation_id, + # "traceId": str(uuid.uuid4()).replace("-", ""), + # "isStartOfSession": 0, + # "productThreadType": "Office", + # "clientInfo": {"clientPlatform": "web"}, "message": { "author": "user", - "inputMethod": "Keyboard", "text": prompt, - "entityAnnotationTypes": ["People", "File", "Event", "Email", "TeamsMessage"], "requestId": str(uuid.uuid4()).replace("-", ""), - "locationInfo": {"timeZoneOffset": 0, "timeZone": "UTC"}, - "locale": "en-US", - "messageType": "Chat", - "experienceType": "Default", }, - "plugins": [], # TODO: support enabling some plugins? + # "plugins": [], } ], - "invocationId": "0", # TODO: should be dynamic? + "invocationId": "0", "target": "chat", "type": CopilotMessageType.USER_PROMPT, }