from cai.sdk.agents import Agent, Runner, gen_trace_id, trace, OpenAIChatCompletionsModel from cai.sdk.agents.mcp import MCPServer, MCPServerSse from cai.sdk.agents.model_settings import ModelSettings #cai tools from cai.tools.reconnaissance.generic_linux_command import generic_linux_command from cai.tools.web.search_web import make_web_search_with_explanation from cai.tools.reconnaissance.exec_code import execute_code from cai.tools.command_and_control.sshpass import run_ssh_command_with_credentials from cai.tools.reconnaissance.shodan import shodan_search, shodan_host_info from cai.tools.web.search_web import make_google_search from cai.tools.misc.reasoning import think from cai.tools.web.google_search import google_search from cai.tools.network.capture_traffic import capture_remote_traffic,remote_capture_session #Portswigger web scraper bot import utils.portswiggerbot as pb #Other utilities import utils.helpers as hp from openai import AsyncOpenAI import asyncio import os import yaml import sys import shutil from datetime import datetime import subprocess import time from typing import Any from dotenv import load_dotenv load_dotenv(override=True) #MAIN PARAMETERS USERNAME = os.getenv("PORTSWIGGER_USERNAME") PASSWORD = os.getenv("PORTSWIGGER_PASSWORD") MODEL = os.getenv("CAI_MODEL") SERVER_URL = "http://127.0.0.1:9876/sse" REMOVE_LABS_WITH_EXPLOIT_SERVER = True # Set to True to remove labs in portswigger that require an additional exploit server to be solved SECTION = "sql-injection" # Change this to the desired section or type of portswigger lab N_LABS = 10 # If you want to test all the labs in the section, change this to -1 AGENT = "webbounty" PROMPT_TYPE = "zero-shot" # Change this to the desired prompt method TOOLS = [generic_linux_command, execute_code] #function to save command terminal logs def setup_tee_logging(log_dir="terminal_output",model="openai/gpt-4o",log_name_prefix="console_log"): """ Save command line outputs of the experiments in folder. Args: log_dir (str): the name of the main directory. log_dir (str): the name of the model used. log_name_prefix (str): the prefix for the log file name. """ model = model.replace("/","-") os.makedirs(log_dir, exist_ok=True) os.makedirs(f"{log_dir}/{model}", exist_ok=True) full_dir = os.path.join(log_dir, model) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") log_path = os.path.join(full_dir, f"{log_name_prefix}_{timestamp}.txt") logfile = open(log_path, "w", buffering=1) tee_out = hp.TeeStream(sys.stdout, logfile) tee_err = hp.TeeStream(sys.stderr, logfile) sys.stdout = tee_out sys.stderr = tee_err return log_path #function to run the agent async def run_agent(agent,user_prompt): """ Run the agent with the provided user prompt. Args: agent (Agent): The CAI Agent to run. user_prompt (str): The user prompt to provide to the agent. Returns: response (class): The response from the agent after processing the user prompt. """ response = await Runner.run(agent, user_prompt) return response def delete_files(folder="logs"): """ Delete all files in the specified folder. Args: folder (str): The folder from which to delete files. Default is "logs". """ # List all files in the folder files = os.listdir(folder) # Check if there are any files if files: for filename in files: file_path = os.path.join(folder, filename) os.remove(file_path) print(f"Deleted file: {filename}") def create_folder_and_move_logs(lab,section,agent,prompt_type,model,lab_status): """ Create a folder structure based on the lab information and move the cai logs to that folder. Args: lab (dict): The lab information containing title, url, and other metadata. section (str): The section or type of lab. agent (str): The name of the agent used. prompt_type (str): The type of prompt used. model (str): The model used for the agent. lab_status (str): The status of the lab (e.g., solved, not-solved, interrupted, unknown). """ #create folder for results model = model.replace("/","-") lab_name = lab['url'].split("/")[-1] if lab_name == 'lab-html-context-nothing-encoded': lab_name = f"{lab_name}-{lab['url'].split("/")[-2]}" if lab_status: lab_status = lab_status.lower().replace(" ","-") destination_folder = os.path.join("results",model, agent, prompt_type, section, lab_status, lab_name) else: destination_folder = os.path.join("results",model, agent, prompt_type, section, lab_name) os.makedirs(destination_folder, exist_ok=True) #move logs folder to results folder logs_file = os.listdir("logs")[0] source_file = os.path.join("logs", logs_file) destination_file = os.path.join(destination_folder, logs_file) shutil.move(source_file, destination_file) async def run(mcp_server: MCPServer): # Create folder log if it doesn't exist os.makedirs("logs", exist_ok=True) #configuration to write terminal output to a file log_file_path = "terminal_output" log_name_prefix = f"{SECTION}.{AGENT}.{PROMPT_TYPE}" _ = setup_tee_logging(log_file_path, model= MODEL, log_name_prefix=log_name_prefix) print(f"[INFO] Console output is being saved to: {log_file_path}") #delete logs files delete_files() #Read Prompt yaml file with open('prompts.yml', 'r') as f: prompts = yaml.load(f, Loader=yaml.SafeLoader) #Create CAI Agent custom_agent = Agent( name="Custom Agent", instructions=prompts['prompts']['system'][AGENT], description="""Agent for AI project to test models with different prompts types on PortSwigger Labs.""", tools=TOOLS, mcp_servers=[mcp_server], model_settings=ModelSettings(tool_choice="required"), model=OpenAIChatCompletionsModel( model=os.getenv('CAI_MODEL', "openai/gpt-4o"), openai_client=AsyncOpenAI(), ) ) print(f"Getting PortSwigger Labs of {SECTION}.....") #Initialize bot bot = pb.Bot(headless=False) #Login to Portswigger Academy bot.login(USERNAME,PASSWORD) #Get labs urls by section or type of lab topics = bot.choose_topic(SECTION) #Get lab metadata #To remove labs that interact with an internal exploit server to be solved if REMOVE_LABS_WITH_EXPLOIT_SERVER: labs = [bot.obtain_lab_information(link) for link in topics] labs = [lab for lab in labs if not 'exploit server' in lab['solution'].lower()] labs = labs[0:N_LABS] else: labs = [bot.obtain_lab_information(link) for link in topics[0:N_LABS]] #format session cookies cookies = bot.driver.get_cookies() # Define names of essential cookies essential_cookies = {'SessionId','Authenticated_UserVerificationId','t'} # Extract only the essential cookies essential_cookies = [cookie for cookie in cookies if cookie['name'] in essential_cookies] print(f"Total Labs extracted: {len(labs)}") #Iterate through each lab for lab in labs: lab_was_interrupted = False # <--- reset for each lab print(f"\n####################################") print(f"Trying to hack Lab: {lab['title']}") user_prompt = prompts['prompts']['user'][PROMPT_TYPE][AGENT] user_prompt = (user_prompt .replace("**LAB_TITLE**", lab['title']) .replace("**LAB_DESCRIPTION**", lab['description']) .replace("**COOKIES**", str(essential_cookies)) .replace("**LAB_ENVIRONMENT_URL**", lab['environment_url']) ) print(f"\nUser Prompt: {user_prompt}") try: response = await run_agent(agent=custom_agent, user_prompt=user_prompt) print("Agent executed successfully.") except KeyboardInterrupt: print("\n[!] User interrupted this lab with Ctrl+C. Skipping lab status check.") response = None lab_was_interrupted = True if lab_was_interrupted: LAB_STATUS = "interrupted" else: try: LAB_STATUS = bot.check_solved_lab(lab['url']) except Exception as e: try: bot.login(USERNAME,PASSWORD) LAB_STATUS = bot.check_solved_lab(lab['url']) except Exception as e: print(f"[!] Error checking lab status: {e}") LAB_STATUS = "unknown" print(f"Lab Status: {LAB_STATUS}") create_folder_and_move_logs(lab, SECTION, AGENT, PROMPT_TYPE, MODEL, LAB_STATUS) #delete logs files and close bot browser delete_files() bot.driver.close() async def main(): async with MCPServerSse( name="SSE Python Server", params={ "url": SERVER_URL, }, ) as server: await run(server) if __name__ == "__main__": asyncio.run(main())