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.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you are still interested in this issue, please respond to keep it open. Thank you!