{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Generic 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", "\n", "To run the Flask app, you also must [download](https://ollama.com/download) and run Ollama, making sure the flask is using a correct model. For example, `ollama run llama3.2:1b` runs the Llama 3.2 1B 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 /workspace/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", "from pyrit.output import output_attack_async\n", "\n", "\n", "def start_flask_app() -> subprocess.Popen:\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", "![image-2.png](../../../assets/playwright_demo.png)\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": {}, "outputs": [], "source": [ "from playwright.async_api import Page, async_playwright\n", "\n", "from pyrit.models import Message\n", "from pyrit.prompt_target import PlaywrightTarget\n", "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", "\n", "await initialize_pyrit_async(memory_db_type=IN_MEMORY) # type: ignore\n", "\n", "\n", "# Define the interaction function\n", "async def interact_with_my_app(page: Page, message: Message) -> 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 (get from first request piece)\n", " prompt_text = message.message_pieces[0].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 `PromptSendingAttack` to send prompts to the target and collects responses." ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", "\u001b[1m\u001b[34m🔹 Turn 1 - USER\u001b[0m\n", "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", "\u001b[34m Tell me a joke about computer programming.\u001b[0m\n", "\n", "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", "\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n", "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", "\u001b[33m Howdy there, partner! I've got one for ya that's fixin' to boot up your funny bone. Here goes:\u001b[0m\n", "\u001b[33m \u001b[0m\n", "\u001b[33m Why did the programmer quit his job?\u001b[0m\n", "\u001b[33m \u001b[0m\n", "\u001b[33m Because he didn't get arrays enough pay!\u001b[0m\n", "\u001b[33m \u001b[0m\n", "\u001b[33m I know, I know, it's a little corny, but I'm a cowboy who code, what can I say?\u001b[0m\n", "\n", "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n" ] } ], "source": [ "import asyncio\n", "import sys\n", "\n", "from pyrit.executor.attack import PromptSendingAttack\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) -> None:\n", " target = PlaywrightTarget(interaction_func=interact_with_my_app, page=page)\n", "\n", " attack = PromptSendingAttack(objective_target=target)\n", "\n", " objective = \"Tell me a joke about computer programming.\"\n", "\n", " result = await attack.execute_async(objective=objective) # type: ignore\n", " await output_attack_async(result)\n", "\n", "\n", "async def run() -> None:\n", " async with async_playwright() as playwright:\n", " browser = await playwright.chromium.launch(headless=True)\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.\")" ] } ], "metadata": { "jupytext": { "main_language": "python" }, "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.13" } }, "nbformat": 4, "nbformat_minor": 5 }