generator: websocket

[leondz]

support websockets (inspired by portswigger's use of them in their LLM tutorial material, https://portswigger.net/web-security/all-labs#web-llm-attacks)

[rc4ne]

@leondz For specific purposes, I think generators would work just fine. Need to be crafted as specific to the task. But I was wondering if there is even a way to generalize it, as we have a JSON template for REST calls since websockets may behave variedly. Example following code worked for portswigger labs btw, attaching for reference correct me if something can be improved in here:

Generator file:

import websocket
import json
import time

def communicate_with_websocket(to_send):
    uri = "wss://<URL>/chat" # WSS PORTSWIGGER LAB URL

    header = {"Cookie":"session=<cookie>"} # COOKIE DURING HANDSHAKE

    ws = websocket.create_connection(uri,headers=header)
    ws.send("READY") # This is required to initiate
    time.sleep(1)
    ws.send(to_send)
    while True:
        response = ws.recv()
        if "Arti Ficial" in response and "content" in response and "CONNECTED" not in response:  # A very very lazy way to catch correct response, as multiple responses are received 
            response = json.loads(response)
            print("\nGOT : " , response)
            return response
            break
        else:
            continue
    ws.close()

def garak_connect(prompt:str, **kwargs)->str:
    input_string = prompt
    input_string = input_string.replace('"', '\\"')
    msg = '{"message":"'+input_string+'"}'         
    response = communicate_with_websocket(msg)
    return response

Invoke file:

import garak
import garak.cli
import generat0r

#USAGE: python invoke_garak.py

garak.cli.main("--model_type function --model_name generat0r#garak_connect --probes promptinject".split())

[leondz]

Thanks very much for this! We would love to support portswigger, and template code is very helpful. The stack is a bit intense right now but I'd like to be able to get to this in the first half of the month.