doc/code/targets/playwright_target.ipynb (288 lines of code) (raw):
{
"cells": [
{
"cell_type": "markdown",
"id": "0",
"metadata": {},
"source": [
"# Playwright Target:\n",
"\n",
"This notebook demonstrates how to interact with the **Playwright Target** in PyRIT.\n",
"\n",
"The `PlaywrightTarget` class allows you to interact with web applications using\n",
"[Playwright](https://playwright.dev/python/docs/intro).\n",
"This is useful for testing interactions with web applications, such as chatbots or other web interfaces,\n",
"especially for red teaming purposes.\n",
"\n",
"## Example Setup\n",
"\n",
"Before you begin, ensure you have the correct version of PyRIT installed and any necessary secrets configured as\n",
"described [here](../../setup/populating_secrets.md).\n",
"\n",
"To run the Flask app, you also must download and run Ollama, making sure the flask is using a correct model.\n",
"\n",
"Additionally, you need to install playwright by executing `playwright install`."
]
},
{
"cell_type": "markdown",
"id": "1",
"metadata": {
"lines_to_next_cell": 0
},
"source": [
"## Example: Interacting with a Web Application using `PlaywrightTarget`\n",
"\n",
"In this example, we'll interact with a simple web application running locally at `http://127.0.0.1:5000`.\n",
"which runs a chatbot that responds to user prompts via ollama\n",
"We'll define an interaction function that navigates to the web application, inputs a prompt, and retrieves the\n",
"bot's response.\n",
"\n",
"## Start the Flask App\n",
"Before we can interact with the web application, we need to start the Flask app that serves the chatbot, this will be done in a subprocess"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Starting Flask app from d:\\git\\PyRIT-internal\\PyRIT\\doc\\code\\targets\\playwright_demo\\app.py...\n",
"Flask app is running and ready!\n"
]
}
],
"source": [
"import os\n",
"import subprocess\n",
"import sys\n",
"import time\n",
"from urllib.error import URLError\n",
"from urllib.request import urlopen\n",
"\n",
"\n",
"def start_flask_app():\n",
" # Get the notebook's directory\n",
" notebook_dir = os.getcwd()\n",
"\n",
" # Construct the path to app.py relative to the notebook directory\n",
" app_path = os.path.join(notebook_dir, \"playwright_demo\", \"app.py\")\n",
"\n",
" # Ensure that app.py exists\n",
" if not os.path.isfile(app_path):\n",
" raise FileNotFoundError(f\"Cannot find app.py at {app_path}\")\n",
"\n",
" # Start the Flask app\n",
" print(f\"Starting Flask app from {app_path}...\")\n",
" flask_process = subprocess.Popen([sys.executable, app_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n",
"\n",
" # Wait for the app to start by checking if the server is up\n",
" server_url = \"http://127.0.0.1:5000\"\n",
" server_ready = False\n",
" while not server_ready:\n",
" try:\n",
" response = urlopen(server_url)\n",
" if response.status == 200:\n",
" server_ready = True\n",
" except URLError:\n",
" time.sleep(0.5) # Wait a bit before checking again\n",
" print(\"Flask app is running and ready!\")\n",
" return flask_process\n",
"\n",
"\n",
"# Start the Flask app\n",
"flask_process = start_flask_app()"
]
},
{
"cell_type": "markdown",
"id": "3",
"metadata": {
"lines_to_next_cell": 0
},
"source": [
"The flask app should now be running locally:\n",
"\n",
"\n",
"\n",
"### Interaction Function\n",
"This is playwright script that interacts with the chatbot web application."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4",
"metadata": {
"lines_to_next_cell": 2
},
"outputs": [],
"source": [
"from playwright.async_api import Page, async_playwright\n",
"\n",
"from pyrit.common import IN_MEMORY, initialize_pyrit\n",
"from pyrit.models import PromptRequestPiece\n",
"from pyrit.orchestrator import PromptSendingOrchestrator\n",
"from pyrit.prompt_target import PlaywrightTarget\n",
"\n",
"initialize_pyrit(memory_db_type=IN_MEMORY)\n",
"\n",
"\n",
"# Define the interaction function\n",
"async def interact_with_my_app(page: Page, request_piece: PromptRequestPiece) -> str:\n",
" # Define selectors\n",
" input_selector = \"#message-input\"\n",
" send_button_selector = \"#send-button\"\n",
" bot_message_selector = \".bot-message\"\n",
"\n",
" # Count existing messages\n",
" initial_messages = await page.query_selector_all(bot_message_selector)\n",
" initial_message_count = len(initial_messages)\n",
"\n",
" # Wait for the page to be ready\n",
" await page.wait_for_selector(input_selector)\n",
"\n",
" # Send the prompt text\n",
" prompt_text = request_piece.converted_value\n",
" await page.fill(input_selector, prompt_text)\n",
" await page.click(send_button_selector)\n",
"\n",
" # Wait for new messages (bot messages)\n",
" await page.wait_for_function(\n",
" f\"document.querySelectorAll('{bot_message_selector}').length > {initial_message_count}\"\n",
" )\n",
"\n",
" # Extract the bot's response text\n",
" bot_message_element = await page.query_selector(f\"{bot_message_selector}:last-child\")\n",
" bot_response = await bot_message_element.text_content()\n",
" return bot_response.strip()"
]
},
{
"cell_type": "markdown",
"id": "5",
"metadata": {
"lines_to_next_cell": 0
},
"source": [
"### Using `PlaywrightTarget` with the Interaction Function and Scorer\n",
"\n",
"Now, we can use the `PlaywrightTarget` by passing the interaction function we defined.\n",
"We'll use the `PromptSendingOrchestrator` to send prompts to the target and collects responses."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6",
"metadata": {
"lines_to_next_cell": 2
},
"outputs": [],
"source": [
"import asyncio\n",
"import sys\n",
"\n",
"if sys.platform == \"win32\":\n",
" asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())\n",
"\n",
"\n",
"# Using PlaywrightTarget with the interaction function and scorer\n",
"async def main(page: Page):\n",
" target = PlaywrightTarget(interaction_func=interact_with_my_app, page=page)\n",
"\n",
" orchestrator = PromptSendingOrchestrator(objective_target=target)\n",
"\n",
" all_prompts = [\n",
" \"Tell me a joke about computer programming.\",\n",
" ]\n",
"\n",
" await orchestrator.send_prompts_async(prompt_list=all_prompts)\n",
" await orchestrator.print_conversations_async() # type: ignore\n",
"\n",
"\n",
"async def run():\n",
" async with async_playwright() as playwright:\n",
" browser = await playwright.chromium.launch(headless=False)\n",
" context = await browser.new_context()\n",
" page: Page = await context.new_page()\n",
" await page.goto(\"http://127.0.0.1:5000\")\n",
" await main(page)\n",
" await context.close()\n",
" await browser.close()\n",
"\n",
"\n",
"# Note in Windows this doesn't run in jupyter notebooks due to playwright limitations\n",
"# https://github.com/microsoft/playwright-python/issues/480\n",
"# await run()\n",
"\n",
"if __name__ == \"__main__\":\n",
" asyncio.run(run())"
]
},
{
"cell_type": "markdown",
"id": "7",
"metadata": {},
"source": [
"## Terminate the Flask App"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Flask app has been terminated.\n"
]
}
],
"source": [
"# Terminate the Flask app when done\n",
"flask_process.terminate()\n",
"flask_process.wait() # Ensure the process has terminated\n",
"print(\"Flask app has been terminated.\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9",
"metadata": {},
"outputs": [],
"source": [
"# Close connection to memory\n",
"from pyrit.memory import CentralMemory\n",
"\n",
"memory = CentralMemory.get_memory_instance()\n",
"memory.dispose_engine()"
]
}
],
"metadata": {
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}