From a8421d2e5a15325523cea86f577dc66ab957cfb5 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Tue, 10 Jun 2025 09:37:33 -0500 Subject: [PATCH 01/12] MultiRest and wsocket Added WebSocket generator and a Rest generator that can execute multiple requests. None should run properly --- garak/generators/multirest.py | 141 ++++++++++++++++++++++++++++++++++ garak/generators/wsocket.py | 57 ++++++++++++++ 2 files changed, 198 insertions(+) create mode 100644 garak/generators/multirest.py create mode 100644 garak/generators/wsocket.py diff --git a/garak/generators/multirest.py b/garak/generators/multirest.py new file mode 100644 index 000000000..a8e7db87d --- /dev/null +++ b/garak/generators/multirest.py @@ -0,0 +1,141 @@ +import json +import logging +import re + +from typing import List, Union +from garak.generators.base import Generator +from garak import _config + +import xml.etree.ElementTree as ET + +import requests + + +DEFAULT_CLASS = "MultiRestGenerator" + + +class MultiRestGenerator(Generator): + DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { + "uri": None, + "template_requests": [], + "template_responses": [], + "burpfile": "file://*.burp", + } + + def __init__(self, config_root=_config): + super().__init__(config_root=config_root) + self.get_reqrep_fromburp(self.burpfile) + + def variable_finder(self, dictionary: dict, locations: dict): + # I apologize for anyone who has to read this code snippet. + key = "" + for (k, v) in dictionary.items(): + if type(v) == str: + for placeholder in re.findall(r'\$[0-9]\$.*\$[0-9]\$', v): + tag = re.search(r'\$[0-9]*\$', placeholder).group(0) + if not locations.get(tag): + locations[tag] = [] + if k == "response_body" or k == "request_body": + locations[tag].append(k) + self.variable_finder(json.loads(v), locations) + #if not locations.get(tag): + # locations[tag] = [k] + else: + locations[tag].append(k) + else: + # this condition currently assumes that the other option is JSON + v_string = json.dumps(v) + for placeholder in re.findall(r'\$[0-9]\$.*\$[0-9]\$', v_string): + tag = re.search(r'\$[0-9]*\$', placeholder).group(0) + if not locations.get(tag): + locations[tag] = [] + if k == "response_body" or k == "request_body": + locations[tag].append(k) + self.variable_finder(v, locations) + else: + locations[tag].append(k) + + def get_reqrep_fromburp(self, burpfile: str): + tree = ET.parse(burpfile) + root = tree.getroot() + for item in root: + for element in item: + if element.tag == "request": + self.template_requests.append( + self.make_reqrep_dictionary(element.text) + ) + if element.tag == "response": + self.template_responses.append( + self.make_reqrep_dictionary(element.text) + ) + + def make_reqrep_dictionary(text: str): + packet = {} + x = text.split("\n") + for y in x: + # change this to regex -> [a-zA-Z]:[a-z\ A-Z]* , something like that + # This condition should parse headers + if ":" in y: + i = y.index(":") + packet[y[:i]] = y[i + 1 :].rstrip("\n").lstrip(" ") + # TODO: This needs to be changed to something more robust + elif " HTTP/" in y: + a = y.rstrip("\n").split(" ") + packet["method"] = a[0] + packet["endpoint"] = a[1] + # TODO: This needs to be changed to something more robust + elif "HTTP/" in y: + a = y.rstrip("\n").split(" ") + packet["status"] = a[1] + packet["error message"] = "".join(a[2:]) + elif not y: + # TODO: There is probably a more robust way of doing this too + if packet.get("endpoint"): + packet["request_body"] = x[-1] + else: + packet["response_body"] = x[-1] + break + return packet + + def grab_value(locations: list, dictionary: dict): + tmp_value = dictionary + for index in locations: + tmp_value = tmp_value[index] + return tmp_value + + def place_value(request_locations, lookup_table, example_request): + tmp_value = example_request + for (k, v) in request_locations.items(): + for index in range(len(v) - 1): + key = v[index] + tmp_value = tmp_value[key] + leaf_key = v[-1] + tmp_value[leaf_key] = lookup_table[k] + + def request_handler(req: dict): + if req["method"] == "GET": + uri = "https://"+ req["Host"] + req["endpoint"] + headers = dict(list(req.items())[2:-1]) + resp = requests.get(uri, headers=headers) + else: + # Assuming POST request + uri = "https://"+ req["Host"] + req["endpoint"] + headers = dict(list(req.items())[2:-1]) + resp = requests.post(uri, headers=headers, json=req["request_body"]) + + def run(self): + output = "" + for i in range(len(self.template_request)) + request_var_locations = self.variable_finder(self.template_request[i]) + self.place_value(request_var_locations, lookup_table, real_request) + + real_response = self.request_handler(real_request) + + response_var_locations = self.variable_finder(self.template_response[i]) + output = self.grab_value(response_var_locations, lookup_table, real_response) + + return output + + + def _call_model(self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: + return ["".join(self.run())] \ No newline at end of file diff --git a/garak/generators/wsocket.py b/garak/generators/wsocket.py new file mode 100644 index 000000000..122c541a3 --- /dev/null +++ b/garak/generators/wsocket.py @@ -0,0 +1,57 @@ +import json +import logging +import re + +from typing import List, Union +from garak.generators.base import Generator +from garak import _config + +import websockets +from websockets.sync.client import connect + +class WebSocketGenerator(Generator): + """ + This is a generator to work with websockets + """ + + DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { + "uri": None, + "auth_key": None, + "temp_request": ["$INPUT"], + "skip": [], + "response_json": False, + "response_json_field": False, + "description": "bruh", + } + + def __init__(self, name="WebSocket", config_root=_config): + # Initialize and validate api_key + super().__init__(name, config_root=config_root) + + def json_handler(self, data): + #TODO: Add try catch + response_json = json.loads(data) + print("AAAAAA") + return json.dumps(response_json[self.response_json_field]) + + + def request(self, request, payload): + with connect(self.uri) as websocket: + websocket.send(request.replace("$INPUT", payload)) + message = websocket.recv() + return self.json_handler(message) if self.response_json == True else message + + def mult_request(self, request, payload): + with connect(self.uri) as websocket: + for requests in self.temp_request: + websocket.send(request.replace("$INPUT", payload)) + + + def _call_model(self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: + return [str(self.request(self.temp_request[0], prompt))] + + + #return [str(self.request(self.temp_request,prompt))] + + +DEFAULT_CLASS = "WebSocketGenerator" \ No newline at end of file From 0a80e02212b56a2a559949b0756b3e0fc70cded4 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Tue, 10 Jun 2025 22:05:09 -0500 Subject: [PATCH 02/12] Update to multirest --- garak/generators/multirest.py | 63 +++++++++++++++++++++++------------ 1 file changed, 42 insertions(+), 21 deletions(-) diff --git a/garak/generators/multirest.py b/garak/generators/multirest.py index a8e7db87d..6e428ef65 100644 --- a/garak/generators/multirest.py +++ b/garak/generators/multirest.py @@ -19,15 +19,15 @@ class MultiRestGenerator(Generator): "uri": None, "template_requests": [], "template_responses": [], - "burpfile": "file://*.burp", + "burpfile": "/Users/host/Downloads/Burp Suite/Burpsuite test", + "description": "Bruh" } - def __init__(self, config_root=_config): - super().__init__(config_root=config_root) + def __init__(self, name="Bruh", config_root=_config): + super().__init__(name, config_root=config_root) self.get_reqrep_fromburp(self.burpfile) def variable_finder(self, dictionary: dict, locations: dict): - # I apologize for anyone who has to read this code snippet. key = "" for (k, v) in dictionary.items(): if type(v) == str: @@ -38,8 +38,6 @@ def variable_finder(self, dictionary: dict, locations: dict): if k == "response_body" or k == "request_body": locations[tag].append(k) self.variable_finder(json.loads(v), locations) - #if not locations.get(tag): - # locations[tag] = [k] else: locations[tag].append(k) else: @@ -69,7 +67,7 @@ def get_reqrep_fromburp(self, burpfile: str): self.make_reqrep_dictionary(element.text) ) - def make_reqrep_dictionary(text: str): + def make_reqrep_dictionary(self, text): packet = {} x = text.split("\n") for y in x: @@ -77,15 +75,15 @@ def make_reqrep_dictionary(text: str): # This condition should parse headers if ":" in y: i = y.index(":") - packet[y[:i]] = y[i + 1 :].rstrip("\n").lstrip(" ") + packet[y[:i]] = y[i + 1 :].rstrip("\n").lstrip(" ").lower() # TODO: This needs to be changed to something more robust elif " HTTP/" in y: - a = y.rstrip("\n").split(" ") + a = y.rstrip('\n').split(' ') packet["method"] = a[0] packet["endpoint"] = a[1] # TODO: This needs to be changed to something more robust elif "HTTP/" in y: - a = y.rstrip("\n").split(" ") + a = y.rstrip('\n').split(' ') packet["status"] = a[1] packet["error message"] = "".join(a[2:]) elif not y: @@ -97,42 +95,65 @@ def make_reqrep_dictionary(text: str): break return packet - def grab_value(locations: list, dictionary: dict): - tmp_value = dictionary - for index in locations: - tmp_value = tmp_value[index] + def grab_value(self, locations: list, dictionary: dict): + tmp_value = None + if dictionary == None: + pass + else: + for k, v in locations.items(): + for index in range(len(v) - 1): + tmp_value = dictionary[v[index]] return tmp_value - def place_value(request_locations, lookup_table, example_request): + def place_value(self, request_locations, lookup_table, example_request): + if not bool(lookup_table): + return tmp_value = example_request for (k, v) in request_locations.items(): for index in range(len(v) - 1): key = v[index] tmp_value = tmp_value[key] leaf_key = v[-1] + print(lookup_table) tmp_value[leaf_key] = lookup_table[k] - def request_handler(req: dict): + def request_handler(self, req: dict): if req["method"] == "GET": uri = "https://"+ req["Host"] + req["endpoint"] headers = dict(list(req.items())[2:-1]) resp = requests.get(uri, headers=headers) + status_code = {"status": str(resp.status_code)} + error_message = {"error message": ""} + headers = dict(resp.headers) + body = {"response_body": resp.text} + response = status_code | error_message | headers | body + return response else: # Assuming POST request uri = "https://"+ req["Host"] + req["endpoint"] headers = dict(list(req.items())[2:-1]) - resp = requests.post(uri, headers=headers, json=req["request_body"]) + resp = requests.post(uri, headers=headers, json=json.loads(req["request_body"])) + status_code = {"status": str(resp.status_code)} + error_message = {"error message": ""} + headers = dict(resp.headers) + body = {"response_body": resp.text} + response = status_code | error_message | headers | body + return response def run(self): + lookup_table = {} output = "" - for i in range(len(self.template_request)) - request_var_locations = self.variable_finder(self.template_request[i]) + request_var_locations = {} + response_var_locations = {} + for i in range(len(self.template_requests)): + real_request = self.template_requests[i].copy() + self.variable_finder(self.template_requests[i], request_var_locations) self.place_value(request_var_locations, lookup_table, real_request) real_response = self.request_handler(real_request) - response_var_locations = self.variable_finder(self.template_response[i]) - output = self.grab_value(response_var_locations, lookup_table, real_response) + self.variable_finder(self.template_responses[i].copy(), response_var_locations) + output = self.grab_value(response_var_locations, real_response) return output From dbd2c5b0d04ccf2d3a49c59646a3718e189e716f Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Wed, 9 Jul 2025 22:00:48 -0500 Subject: [PATCH 03/12] WebSocketGenerator --- garak/generators/wsocket.py | 32 ++++++++------------- tests/generators/test_websocket.py | 45 ++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 21 deletions(-) create mode 100644 tests/generators/test_websocket.py diff --git a/garak/generators/wsocket.py b/garak/generators/wsocket.py index 122c541a3..1011281aa 100644 --- a/garak/generators/wsocket.py +++ b/garak/generators/wsocket.py @@ -17,41 +17,31 @@ class WebSocketGenerator(Generator): DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { "uri": None, "auth_key": None, - "temp_request": ["$INPUT"], - "skip": [], - "response_json": False, - "response_json_field": False, - "description": "bruh", + "body": '{}', + "json_response": True, + "json_key": "output", } def __init__(self, name="WebSocket", config_root=_config): - # Initialize and validate api_key super().__init__(name, config_root=config_root) def json_handler(self, data): - #TODO: Add try catch response_json = json.loads(data) - print("AAAAAA") - return json.dumps(response_json[self.response_json_field]) + return json.dumps(response_json[self.json_key]) - def request(self, request, payload): + def request(self, payload): with connect(self.uri) as websocket: - websocket.send(request.replace("$INPUT", payload)) + websocket.send(self.body.replace("$INPUT", payload)) message = websocket.recv() - return self.json_handler(message) if self.response_json == True else message - - def mult_request(self, request, payload): - with connect(self.uri) as websocket: - for requests in self.temp_request: - websocket.send(request.replace("$INPUT", payload)) + return self.json_handler(message) if self.json_response == True else message def _call_model(self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: - return [str(self.request(self.temp_request[0], prompt))] - - - #return [str(self.request(self.temp_request,prompt))] + if output := self.request(self, prompt) == dict : + return output[self.json_key] + else : + return output DEFAULT_CLASS = "WebSocketGenerator" \ No newline at end of file diff --git a/tests/generators/test_websocket.py b/tests/generators/test_websocket.py new file mode 100644 index 000000000..1f02249e3 --- /dev/null +++ b/tests/generators/test_websocket.py @@ -0,0 +1,45 @@ +import pytest +from unittest.mock import patch, MagicMock + +from garak.generators.wsocket import WebSocketGenerator + +@pytest.fixture +def ws_gen(): + gen = WebSocketGenerator() + gen.uri = "ws://test" + gen.body = '{"input": "$INPUT"}' + gen.json_response = True + gen.json_key = "output" + return gen + +def test_json_handler(ws_gen): + data = '{"output": "test response"}' + result = ws_gen.json_handler(data) + assert result == '"test response"' + +@patch("garak.generators.wsocket.connect") +def test_request_json_response(mock_connect, ws_gen): + mock_ws = MagicMock() + mock_ws.recv.return_value = '{"output": "foo"}' + mock_connect.return_value.__enter__.return_value = mock_ws + + result = ws_gen.request("bar") + assert result == '"foo"' + mock_ws.send.assert_called_once_with('{"input": "bar"}') + +@patch("garak.generators.wsocket.connect") +def test_request_raw_response(mock_connect, ws_gen): + ws_gen.json_response = False + mock_ws = MagicMock() + mock_ws.recv.return_value = "raw" + mock_connect.return_value.__enter__.return_value = mock_ws + + result = ws_gen.request("baz") + assert result == "raw" + +def test_live_request_raw_response(ws_gen): + ws_gen.json_response = False + ws_gen.uri = "wss://echo.websocket.events" + result = ws_gen.request("test") + assert result == "test" + \ No newline at end of file From 7fe4d052917c9eff573977d8256ff0a7aa152064 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Thu, 10 Jul 2025 09:55:20 -0500 Subject: [PATCH 04/12] MultiRestGenerator --- garak/generators/multirest.py | 347 ++++++++++++++++++----------- tests/generators/test_multirest.py | 104 +++++++++ tools/rest/multirest.xml | 104 +++++++++ 3 files changed, 431 insertions(+), 124 deletions(-) create mode 100644 tests/generators/test_multirest.py create mode 100644 tools/rest/multirest.xml diff --git a/garak/generators/multirest.py b/garak/generators/multirest.py index 6e428ef65..8680f11b8 100644 --- a/garak/generators/multirest.py +++ b/garak/generators/multirest.py @@ -1,5 +1,4 @@ import json -import logging import re from typing import List, Union @@ -15,148 +14,248 @@ class MultiRestGenerator(Generator): + """ + This is a generator to work with multiple request-response pairs using Burp Suite's exported XML format. It supports variable substitution, in case you need to build succeeding requests uponinformation from previous request/response pairs. + The generator can handle both JSON and non-JSON responses. + + To set a variable in the response, use the format $n$...$n$, where n is a number. + + For example: + HTTP/2 200 OK + Date: Tue, 20 May 2025 05:58:32 GMT + Content-Type: application/json + Content-Length: 52 + Server: uvicorn + Apigw-Requestid: $2$K2kBRiewPHcEM8A=$2$ + Access-Control-Allow-Origin: https://platform.dreadnode.io + + {"id":"$1$999999$1$"}]]> + + The above response will set the variable 1 to "999999" and variable 2 to "K2kBRiewPHcEM8A=" + + From there you can place these variables in the request body or headers like this: + GET /score?id=$1$123$1$ HTTP/2 + Host: example.com + User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:138.0) Gecko/20100101 Firefox/138.0 + Accept: */* + Accept-Language: en-US,en;q=0.5 + Accept-Encoding: gzip, deflate, br + Referer: https://example.com/score + Content-Type: application/json + Origin: https://example.com + Sec-Fetch-Dest: empty + Sec-Fetch-Mode: cors + Sec-Fetch-Site: same-site + Dnt: 1 + Sec-Gpc: 1 + Priority: u=0 + Te: trailers + + The above request will replace $1$123$1$ with the value of variable 1, which is "999999". + + You can specify where the $INPUT and $OUTPUT variables go in the request and response, as well: + + POST /score HTTP/2 + Host: example.com + User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:138.0) Gecko/20100101 Firefox/138.0 + Accept: */* + Accept-Language: en-US,en;q=0.5 + Accept-Encoding: gzip, deflate, br + Content-Type: application/json + Sec-Fetch-Mode: cors + Sec-Fetch-Site: same-site + Dnt: 1 + Sec-Gpc: 1 + Priority: u=0 + Te: trailers + + {"data":"$INPUT"} + + Here $INPUT will be replaced with the input value provided to call_model. + + HTTP/2 400 Bad Request + Date: Tue, 20 May 2025 05:58:14 GMT + Content-Type: application/json + Content-Length: 46 + Server: uvicorn + Apigw-Requestid: K2j-dj6KPHcEMvw= + Access-Control-Allow-Origin: https://platform.dreadnode.io + + {"message":"$OUTPUT"}]]> + + Here the location of $OUTPUT in the response body indicates that the output value will be extracted from the response JSON under the "message" key. + If the response is not JSON, it will be returned as is. + """ DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { - "uri": None, - "template_requests": [], - "template_responses": [], - "burpfile": "/Users/host/Downloads/Burp Suite/Burpsuite test", - "description": "Bruh" + "reqresp_pairs": [], + "variables": {}, + "burpfile": "./tools/rest/multirest.xml", } - def __init__(self, name="Bruh", config_root=_config): + def __init__(self, name="MultiRestGenerator", config_root=_config): super().__init__(name, config_root=config_root) - self.get_reqrep_fromburp(self.burpfile) - - def variable_finder(self, dictionary: dict, locations: dict): - key = "" - for (k, v) in dictionary.items(): - if type(v) == str: - for placeholder in re.findall(r'\$[0-9]\$.*\$[0-9]\$', v): - tag = re.search(r'\$[0-9]*\$', placeholder).group(0) - if not locations.get(tag): - locations[tag] = [] - if k == "response_body" or k == "request_body": - locations[tag].append(k) - self.variable_finder(json.loads(v), locations) - else: - locations[tag].append(k) - else: - # this condition currently assumes that the other option is JSON - v_string = json.dumps(v) - for placeholder in re.findall(r'\$[0-9]\$.*\$[0-9]\$', v_string): - tag = re.search(r'\$[0-9]*\$', placeholder).group(0) - if not locations.get(tag): - locations[tag] = [] - if k == "response_body" or k == "request_body": - locations[tag].append(k) - self.variable_finder(v, locations) - else: - locations[tag].append(k) def get_reqrep_fromburp(self, burpfile: str): tree = ET.parse(burpfile) root = tree.getroot() - for item in root: - for element in item: - if element.tag == "request": - self.template_requests.append( - self.make_reqrep_dictionary(element.text) - ) - if element.tag == "response": - self.template_responses.append( - self.make_reqrep_dictionary(element.text) - ) + pairs = [] + for item in root.findall("item"): + req = item.find("request").text + resp = item.find("response").text + pairs.append( + { + "request": self.make_reqrep_dictionary(req), + "response": self.make_reqrep_dictionary(resp), + } + ) + return pairs def make_reqrep_dictionary(self, text): - packet = {} - x = text.split("\n") - for y in x: + packet = {"headers": {}, "body": ""} + http_line = text.split("\n") + for substring in http_line: # change this to regex -> [a-zA-Z]:[a-z\ A-Z]* , something like that # This condition should parse headers - if ":" in y: - i = y.index(":") - packet[y[:i]] = y[i + 1 :].rstrip("\n").lstrip(" ").lower() + if ":" in substring: + i = substring.index(":") + packet["headers"][substring[:i]] = ( + substring[i + 1 :].rstrip("\n").lstrip(" ").lower() + ) # TODO: This needs to be changed to something more robust - elif " HTTP/" in y: - a = y.rstrip('\n').split(' ') + elif " HTTP/" in substring: + a = substring.rstrip("\n").split(" ") packet["method"] = a[0] packet["endpoint"] = a[1] # TODO: This needs to be changed to something more robust - elif "HTTP/" in y: - a = y.rstrip('\n').split(' ') + elif "HTTP/" in substring: + a = substring.rstrip("\n").split(" ") packet["status"] = a[1] packet["error message"] = "".join(a[2:]) - elif not y: + elif not substring: # TODO: There is probably a more robust way of doing this too - if packet.get("endpoint"): - packet["request_body"] = x[-1] - else: - packet["response_body"] = x[-1] - break + packet["body"] = http_line[-1] return packet - def grab_value(self, locations: list, dictionary: dict): - tmp_value = None - if dictionary == None: - pass - else: - for k, v in locations.items(): - for index in range(len(v) - 1): - tmp_value = dictionary[v[index]] - return tmp_value - - def place_value(self, request_locations, lookup_table, example_request): - if not bool(lookup_table): - return - tmp_value = example_request - for (k, v) in request_locations.items(): - for index in range(len(v) - 1): - key = v[index] - tmp_value = tmp_value[key] - leaf_key = v[-1] - print(lookup_table) - tmp_value[leaf_key] = lookup_table[k] - - def request_handler(self, req: dict): - if req["method"] == "GET": - uri = "https://"+ req["Host"] + req["endpoint"] - headers = dict(list(req.items())[2:-1]) - resp = requests.get(uri, headers=headers) - status_code = {"status": str(resp.status_code)} - error_message = {"error message": ""} - headers = dict(resp.headers) - body = {"response_body": resp.text} - response = status_code | error_message | headers | body - return response + def extract_placeholders(self, text): + # Finds all $n$...$n$ placeholders and returns a dict {n: value} + matches = re.findall(r"\$(\d+)\$(.*?)\$\1\$", text) + return {num: val for num, val in matches} + + def substitute_placeholders(self, text, input_value=None): + # Replace $n$...$n$ with self.variables[n], and $INPUT with input_value + def repl(match): + num = match.group(1) + return self.variables.get(num, match.group(0)) + + text = re.sub(r"\$(\d+)\$.*?\$\1\$", repl, text) + if input_value is not None: + text = text.replace("$INPUT", input_value) + return text + + def substitute_in_packet(self, packet, input_value=None): + # Substitute placeholders in headers and body + new_packet = packet.copy() + new_packet["headers"] = { + k: self.substitute_placeholders(v, input_value) + for k, v in packet["headers"].items() + } + new_packet["body"] = self.substitute_placeholders( + packet.get("body", ""), input_value + ) + new_packet["endpoint"] = self.substitute_placeholders( + packet.get("endpoint", ""), input_value + ) + return new_packet + + def extract_vars_from_packet(self, packet): + # Extract variables from headers and body + vars_found = {} + for v in packet["headers"].values(): + vars_found.update(self.extract_placeholders(v)) + vars_found.update(self.extract_placeholders(packet.get("body", ""))) + return vars_found + + def make_request(self, packet): + method = packet.get("method", "POST") + url = f"https://{packet['headers'].get('Host')}{packet.get('endpoint', '')}" + headers = { + k: v + for k, v in packet["headers"].items() + if k not in ["method", "Host", "endpoint"] + } + data = packet.get("body", "") + if method == "GET": + resp = requests.get(url, headers=headers) else: - # Assuming POST request - uri = "https://"+ req["Host"] + req["endpoint"] - headers = dict(list(req.items())[2:-1]) - resp = requests.post(uri, headers=headers, json=json.loads(req["request_body"])) - status_code = {"status": str(resp.status_code)} - error_message = {"error message": ""} - headers = dict(resp.headers) - body = {"response_body": resp.text} - response = status_code | error_message | headers | body - return response - - def run(self): - lookup_table = {} - output = "" - request_var_locations = {} - response_var_locations = {} - for i in range(len(self.template_requests)): - real_request = self.template_requests[i].copy() - self.variable_finder(self.template_requests[i], request_var_locations) - self.place_value(request_var_locations, lookup_table, real_request) - - real_response = self.request_handler(real_request) - - self.variable_finder(self.template_responses[i].copy(), response_var_locations) - output = self.grab_value(response_var_locations, real_response) - - return output - - - def _call_model(self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: - return ["".join(self.run())] \ No newline at end of file + try: + json_data = json.loads(data) + resp = requests.post(url, headers=headers, json=json_data) + except Exception: + resp = requests.post(url, headers=headers, data=data) + return resp + + def compare_responses(self, expected, actual): + if "$OUTPUT" in expected: + return True + try: + expected_json = json.loads(expected) + actual_json = actual.json() + return expected_json == actual_json + except Exception: + return expected.strip() == actual.text.strip() + + def extract_output_from_json(self, template, real): + outputs = [] + + if isinstance(template, dict) and isinstance(real, dict): + for k, v in template.items(): + if v == "$OUTPUT": + outputs.append(real.get(k)) + elif isinstance(v, (dict, list)): + outputs.extend(self.extract_output_from_json(v, real.get(k, {}))) + elif isinstance(template, list) and isinstance(real, list): + for t_item, r_item in zip(template, real): + outputs.extend(self.extract_output_from_json(t_item, r_item)) + return outputs + + def run(self, prompt): + for pair in self.reqresp_pairs: + # Substitute variables and $INPUT in request headers and body + req_packet = self.substitute_in_packet(pair["request"], prompt) + + # Make the request + resp = self.make_request(req_packet) + + # Extract variables from expected response headers and body + vars_from_resp = self.extract_vars_from_packet(pair["response"]) + self.variables.update(vars_from_resp) + + # Compare actual and expected response + expected_body = pair["response"].get("body", "") + real_body = resp.text + + expected_body_sub = self.substitute_placeholders( + pair["response"].get("body", "") + ) + if self.compare_responses(expected_body_sub, resp): + # If $OUTPUT is in the expected body and both are JSON, extract output + if "$OUTPUT" in expected_body: + try: + expected_json = json.loads(expected_body) + real_json = resp.json() + outputs = self.extract_output_from_json(expected_json, real_json) + output = None + for key in range(len(outputs)): + output = outputs[key] + return output + except Exception: + return resp.text + else: + raise ValueError( + f"Response mismatch: expected {expected_body_sub}, got {real_body}" + ) + + def _call_model( + self, prompt: str, generations_this_call: int = 1 + ) -> List[Union[str, None]]: + return ["".join(self.run(prompt))] diff --git a/tests/generators/test_multirest.py b/tests/generators/test_multirest.py new file mode 100644 index 000000000..35e39923c --- /dev/null +++ b/tests/generators/test_multirest.py @@ -0,0 +1,104 @@ +import pytest +from garak.generators.multirest import MultiRestGenerator + +BURP_XML_PATH = "./tools/rest/multirest.xml" + +@pytest.fixture +def mr_gen(tmp_path): + generator = MultiRestGenerator() + generator.burpfile = BURP_XML_PATH + return generator + +def test_extract_output_from_json(): + gen = MultiRestGenerator() + template = {"message": "$OUTPUT"} + real = {"message": "error details"} + outputs = gen.extract_output_from_json(template, real) + assert outputs == ["error details"] + +def test_run_extracts_output(mr_gen, requests_mock): + # Prepare the generator with the test XML + mr_gen.reqresp_pairs = mr_gen.get_reqrep_fromburp(BURP_XML_PATH) + + # Mock the endpoint for both test cases + url = "https://example.com/score" + # First response: 200 OK + requests_mock.post(url, [ + { + "status_code": 200, + "json": {"id": "999999"} + }, + ]) + url2 = "https://example.com/score?id=999999" + requests_mock.get(url2, status_code=400, json={"message": "error details"}) + + # Run the generator for the first request (should not extract $OUTPUT) + output = mr_gen.run("Hello") + # The second request in the XML expects $OUTPUT to be extracted + assert output == "error details" + + +def test_substitute_placeholders(mr_gen): + mr_gen = MultiRestGenerator() + mr_gen.variables = {"1": "foo"} + text = '{"data":"$1$Hello$1$"}' + result = mr_gen.substitute_placeholders(text) + assert result == '{"data":"foo"}' + +def test_response_mismatch_raises(mr_gen, requests_mock): + mr_gen.reqresp_pairs = [{ + "request": { + "method": "GET", + "headers": {"Host": "example.com"}, + "endpoint": "/fail", + "body": "" + }, + "response": { + "body": '{"expected":"value"}', + "headers": {} + } + }] + url = "https://example.com/fail" + requests_mock.get(url, json={"unexpected": "different"}) + with pytest.raises(ValueError): + mr_gen.run("prompt") + +def test_substitute_placeholders_no_vars(mr_gen): + mr_gen.variables = {} + text = '{"data":"$1$Hello$1$"}' + result = mr_gen.substitute_placeholders(text) + assert result == '{"data":"$1$Hello$1$"}' + +def test_extract_vars_from_headers(mr_gen): + packet = { + "headers": {"X-Token": "$1$tokenval$1$"}, + "body": "" + } + vars_found = mr_gen.extract_vars_from_packet(packet) + assert vars_found == {"1": "tokenval"} + +def test_non_json_response(mr_gen, requests_mock): + packet = { + "method": "GET", + "headers": {"Host": "example.com"}, + "endpoint": "/plain", + "body": "" + } + url = "https://example.com/plain" + requests_mock.get(url, text="plain text response") + resp = mr_gen.make_request(packet) + assert resp.text == "plain text response" + +def test_get_request_with_placeholder(mr_gen, requests_mock): + mr_gen.variables = {"1": "abc123"} + packet = { + "method": "GET", + "headers": {"Host": "example.com"}, + "endpoint": "/score?id=$1$foo$1$", + "body": "" + } + url = "https://example.com/score?id=abc123" + requests_mock.get(url, text="ok") + req_packet = mr_gen.substitute_in_packet(packet) + resp = mr_gen.make_request(req_packet) + assert resp.text == "ok" \ No newline at end of file diff --git a/tools/rest/multirest.xml b/tools/rest/multirest.xml new file mode 100644 index 000000000..66987fee2 --- /dev/null +++ b/tools/rest/multirest.xml @@ -0,0 +1,104 @@ + + + + + + + + + + + + + + + + + + + + + + + +]> + + + + + example.com + 443 + https + + + null + + 200 + 270 + JSON + + + + + + + example.com + 443 + https + + + null + + 400 + 273 + JSON + + + + From de8c25fd93f00e5cb9c1cd941a8567f14c99f7f1 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Fri, 18 Jul 2025 22:13:46 -0500 Subject: [PATCH 05/12] Remake of the MultiRestGenerator --- garak/data/rest/multirest.xml | 104 ++++++++++ garak/generators/multirest.py | 318 ++++++++++++++++++----------- pyproject.toml | 1 + requirements.txt | 1 + tests/generators/test_multirest.py | 154 ++++++-------- 5 files changed, 376 insertions(+), 202 deletions(-) create mode 100644 garak/data/rest/multirest.xml diff --git a/garak/data/rest/multirest.xml b/garak/data/rest/multirest.xml new file mode 100644 index 000000000..645ceaec7 --- /dev/null +++ b/garak/data/rest/multirest.xml @@ -0,0 +1,104 @@ + + + + + + + + + + + + + + + + + + + + + + + +]> + + + + + example.com + 443 + https + + + null + + 200 + 270 + JSON + + + + + + + example.com + 443 + https + + + null + + 400 + 273 + JSON + + + + diff --git a/garak/generators/multirest.py b/garak/generators/multirest.py index 8680f11b8..979f7ef75 100644 --- a/garak/generators/multirest.py +++ b/garak/generators/multirest.py @@ -14,13 +14,15 @@ class MultiRestGenerator(Generator): - """ - This is a generator to work with multiple request-response pairs using Burp Suite's exported XML format. It supports variable substitution, in case you need to build succeeding requests uponinformation from previous request/response pairs. - The generator can handle both JSON and non-JSON responses. + """ + This is a generator to work with multiple request-response pairs using Burp Suite's exported XML format. It supports variable substitution, in case you need to build succeeding requests uponinformation from previous request/response pairs. + The generator supports only JSON request and responses. To set a variable in the response, use the format $n$...$n$, where n is a number. - For example: + When you save a Burpfile with reqeust and response items, be sure to uncheck the base64 button. + + For example: HTTP/2 200 OK Date: Tue, 20 May 2025 05:58:32 GMT Content-Type: application/json @@ -43,7 +45,7 @@ class MultiRestGenerator(Generator): Referer: https://example.com/score Content-Type: application/json Origin: https://example.com - Sec-Fetch-Dest: empty + Sec-Fetch-Dest: empty Sec-Fetch-Mode: cors Sec-Fetch-Site: same-site Dnt: 1 @@ -70,7 +72,7 @@ class MultiRestGenerator(Generator): Te: trailers {"data":"$INPUT"} - + Here $INPUT will be replaced with the input value provided to call_model. HTTP/2 400 Bad Request @@ -80,23 +82,107 @@ class MultiRestGenerator(Generator): Server: uvicorn Apigw-Requestid: K2j-dj6KPHcEMvw= Access-Control-Allow-Origin: https://platform.dreadnode.io - + {"message":"$OUTPUT"}]]> Here the location of $OUTPUT in the response body indicates that the output value will be extracted from the response JSON under the "message" key. If the response is not JSON, it will be returned as is. + + You can also specify request-response pairs in a config file similar to that of the Rest Generator: + { + "multi_endpoint_rest": { + "MultiEndpointGenerator": { + "first_stage": { + "name": "request service", + "uri": "https://example.ai/llm", + "method": "post", + "headers": { + "X-Authorization": "$KEY", + }, + "req_template_json_object": { + "text": "$INPUT" + }, + "response_json": true, + "response_json_body": "$1$Hello$1$" + }, + "second_stage": { + "name": "response service", + "uri": "https://example.ai/llm", + "method": "post", + "headers": { + "X-Authorization": "$KEY", + }, + "req_template_json_object": { + "text": "$INPUT" + }, + "response_json": true, + "response_template_json_object": { + "text": "$OUT$data$OUT$" + } + } + } + } + } """ + DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { "reqresp_pairs": [], - "variables": {}, - "burpfile": "./tools/rest/multirest.xml", + "burpfile": None, + "manual": [], } def __init__(self, name="MultiRestGenerator", config_root=_config): super().__init__(name, config_root=config_root) + # Error condition for when neither burpfile nor manual is provided + if not self.burpfile and not self.manual: + raise ValueError( + "You must provide either a burpfile or manually define request and responses in the config file for MultiRestGenerator." + ) + # Error condition for when both burpfile and manual is provided + if self.burpfile and self.manual: + raise ValueError( + "You cannot provide both a burpfile and manual definitions to the MultiRestGenerator." + ) + # Load request-response pairs from burpfile or manual + if self.burpfile: + self.reqresp_pairs = self.get_reqrep_fromburp() + else: + self.reqresp_pairs = self.get_reqrep_fromconfig() + + def get_reqrep_fromconfig(self): + def parse_domain(uri): + # Extracts the domain from a URI + match = re.match(r"https?://([^/]+)", uri) + return match.group(1) if match else None - def get_reqrep_fromburp(self, burpfile: str): - tree = ET.parse(burpfile) + # Parses the manual definitions and returns a list of request-response pairs + pairs = [] + for k, item in self.manual.items(): + # Might need to just move this to a separate function. + req = { + "headers": item["headers"] | {"Host": parse_domain(item["uri"])}, + "body": ( + item["req_template_json_object"] if item["method"] == "post" else {} + ), + "method": item["method"], + "endpoint": item["uri"][item["uri"].rindex("/") :], + } + resp = { + "headers": {}, + "body": ( + item["response_template_json_object"] + if item.get("response_json") + else {} + ), + "status": "", + "error message": "", + } + pairs.append({"request": req, "response": resp}) + return pairs + + def get_reqrep_fromburp(self): + # Parses the Burp Suite XML file and extracts request-response pairs + tree = ET.parse(self.burpfile) root = tree.getroot() pairs = [] for item in root.findall("item"): @@ -110,71 +196,34 @@ def get_reqrep_fromburp(self, burpfile: str): ) return pairs - def make_reqrep_dictionary(self, text): + def make_reqrep_dictionary(self, text: str): + # Converts a raw HTTP request or response text into a dictionary format packet = {"headers": {}, "body": ""} http_line = text.split("\n") - for substring in http_line: - # change this to regex -> [a-zA-Z]:[a-z\ A-Z]* , something like that - # This condition should parse headers + border = http_line.index("") + for substring in http_line[:border]: + # This condition assumes that the line is a Header if ":" in substring: i = substring.index(":") packet["headers"][substring[:i]] = ( substring[i + 1 :].rstrip("\n").lstrip(" ").lower() ) - # TODO: This needs to be changed to something more robust + # This condition assumes that the line is a Request elif " HTTP/" in substring: a = substring.rstrip("\n").split(" ") packet["method"] = a[0] packet["endpoint"] = a[1] - # TODO: This needs to be changed to something more robust + # This condition assumes that the line is a Response elif "HTTP/" in substring: a = substring.rstrip("\n").split(" ") packet["status"] = a[1] packet["error message"] = "".join(a[2:]) - elif not substring: - # TODO: There is probably a more robust way of doing this too - packet["body"] = http_line[-1] + try: + packet["body"] = json.loads(http_line[-1]) + except json.decoder.JSONDecodeError: + packet["body"] = http_line[-1] return packet - def extract_placeholders(self, text): - # Finds all $n$...$n$ placeholders and returns a dict {n: value} - matches = re.findall(r"\$(\d+)\$(.*?)\$\1\$", text) - return {num: val for num, val in matches} - - def substitute_placeholders(self, text, input_value=None): - # Replace $n$...$n$ with self.variables[n], and $INPUT with input_value - def repl(match): - num = match.group(1) - return self.variables.get(num, match.group(0)) - - text = re.sub(r"\$(\d+)\$.*?\$\1\$", repl, text) - if input_value is not None: - text = text.replace("$INPUT", input_value) - return text - - def substitute_in_packet(self, packet, input_value=None): - # Substitute placeholders in headers and body - new_packet = packet.copy() - new_packet["headers"] = { - k: self.substitute_placeholders(v, input_value) - for k, v in packet["headers"].items() - } - new_packet["body"] = self.substitute_placeholders( - packet.get("body", ""), input_value - ) - new_packet["endpoint"] = self.substitute_placeholders( - packet.get("endpoint", ""), input_value - ) - return new_packet - - def extract_vars_from_packet(self, packet): - # Extract variables from headers and body - vars_found = {} - for v in packet["headers"].values(): - vars_found.update(self.extract_placeholders(v)) - vars_found.update(self.extract_placeholders(packet.get("body", ""))) - return vars_found - def make_request(self, packet): method = packet.get("method", "POST") url = f"https://{packet['headers'].get('Host')}{packet.get('endpoint', '')}" @@ -192,68 +241,107 @@ def make_request(self, packet): resp = requests.post(url, headers=headers, json=json_data) except Exception: resp = requests.post(url, headers=headers, data=data) - return resp + return self.parse_response(resp) def compare_responses(self, expected, actual): - if "$OUTPUT" in expected: - return True - try: - expected_json = json.loads(expected) - actual_json = actual.json() - return expected_json == actual_json - except Exception: - return expected.strip() == actual.text.strip() - - def extract_output_from_json(self, template, real): - outputs = [] - - if isinstance(template, dict) and isinstance(real, dict): - for k, v in template.items(): - if v == "$OUTPUT": - outputs.append(real.get(k)) - elif isinstance(v, (dict, list)): - outputs.extend(self.extract_output_from_json(v, real.get(k, {}))) - elif isinstance(template, list) and isinstance(real, list): - for t_item, r_item in zip(template, real): - outputs.extend(self.extract_output_from_json(t_item, r_item)) - return outputs - - def run(self, prompt): - for pair in self.reqresp_pairs: - # Substitute variables and $INPUT in request headers and body - req_packet = self.substitute_in_packet(pair["request"], prompt) - - # Make the request - resp = self.make_request(req_packet) - - # Extract variables from expected response headers and body - vars_from_resp = self.extract_vars_from_packet(pair["response"]) - self.variables.update(vars_from_resp) - - # Compare actual and expected response - expected_body = pair["response"].get("body", "") - real_body = resp.text + # This just checks if the keys match. + e_keys = list(expected.keys()).sort() + a_keys = list(actual.keys()).sort() + return e_keys == a_keys + + def extract_variable_locations(self, expected, location=None, locations=None): + if location == None: + location = [] + if locations == None: + locations = {} + + if type(expected) == dict: + for k, v in expected.items(): + self.extract_variable_locations(v, location + [k], locations) + else: + placeholders = re.findall(r"\$(OUT|\d+)\$", expected) + for placeholder in placeholders: + locations[placeholder] = location + return locations + + def update_locations_dictionary(self, dictionary, packet): + location_dict = dictionary.copy() + result = packet + for var_number, locations in location_dict.items(): + if type(locations) == list: + for location in locations: + result = result[location] + else: + continue + location_dict[var_number] = result + return location_dict + + def place_var_into_request_packet(self, packet, locations, new_packet=None): + if not new_packet: + new_packet = {} + if type(packet) == dict: + for k, v in packet.items(): + new_packet[k] = self.place_var_into_request_packet(v, locations) + return new_packet + else: + string_to_edit = packet # At this point, this should be a string. + + def repl(match): + num = match.group(1) + return locations.get(num, match.group(0)) + + return re.sub(r"\$(\d+)\$.*?\$\1\$", repl, string_to_edit) + + def place_input_into_request_packet( + self, packet: dict, prompt: str, new_packet=None + ): + if not new_packet: + new_packet = {} + if type(packet) == dict: + for k, v in packet.items(): + new_packet[k] = self.place_input_into_request_packet(v, prompt) + return new_packet + else: + string_to_compare = packet # At this point, this should be a string. + if string_to_compare == "$INPUT": + return prompt + else: + return string_to_compare + + def parse_response(self, raw_response: requests.Response) -> dict: + return { + "headers": raw_response.headers, + "body": ( + raw_response.json() + if raw_response.headers["Content-Type"] == "application/json" + else raw_response.text + ), + } - expected_body_sub = self.substitute_placeholders( - pair["response"].get("body", "") - ) - if self.compare_responses(expected_body_sub, resp): - # If $OUTPUT is in the expected body and both are JSON, extract output - if "$OUTPUT" in expected_body: - try: - expected_json = json.loads(expected_body) - real_json = resp.json() - outputs = self.extract_output_from_json(expected_json, real_json) - output = None - for key in range(len(outputs)): - output = outputs[key] - return output - except Exception: - return resp.text + def run(self, prompt: str,locations = {}) -> str: + + for pair in self.reqresp_pairs: + req_packet = self.place_input_into_request_packet(pair["request"], prompt) + req_packet = self.place_var_into_request_packet(req_packet, locations) + resp_packet = self.make_request(req_packet) + if self.compare_responses(pair["response"], resp_packet): + try: + locations.update(self.extract_variable_locations(pair["response"])) + locations.update( + self.update_locations_dictionary(locations, resp_packet) + ) + except: + raise KeyError( + f"Locator failure: Variable could not be found in actual response: {locations} -> {resp_packet}" + ) else: raise ValueError( - f"Response mismatch: expected {expected_body_sub}, got {real_body}" + f"Response mismatch: expected {pair["response"]}, got {resp_packet}" ) + if locations.get("OUT"): + return locations["OUT"] + else: + return "" def _call_model( self, prompt: str, generations_this_call: int = 1 diff --git a/pyproject.toml b/pyproject.toml index eccd8a931..e89744ee0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,6 +74,7 @@ dependencies = [ "numpy>=1.26.1", "zalgolib>=0.2.2", "ecoji>=0.1.1", + "websockets>=14.0", "deepl==1.17.0", "fschat>=0.2.36", "litellm>=1.41.21", diff --git a/requirements.txt b/requirements.txt index 50de30fe5..1f3a7ed44 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,6 +24,7 @@ markdown>=3.4.3 numpy>=1.26.1 zalgolib>=0.2.2 ecoji>=0.1.1 +websockets>=14.0 deepl==1.17.0 fschat>=0.2.36 litellm>=1.41.21 diff --git a/tests/generators/test_multirest.py b/tests/generators/test_multirest.py index 35e39923c..a68b3a7dd 100644 --- a/tests/generators/test_multirest.py +++ b/tests/generators/test_multirest.py @@ -1,104 +1,84 @@ import pytest from garak.generators.multirest import MultiRestGenerator +from garak import _config, _plugins -BURP_XML_PATH = "./tools/rest/multirest.xml" +EXAMPLE_BURP_XML = { + "burpfile":"./garak/data/rest/multirest.xml" +} + +EXAMPLE_CONFIG = { + "multirest": { + "MultiRestGenerator": { + "manual": { + "first_stage": { + "name": "request service", + "uri": "https://example.ai/llm", + "method": "POST", + "headers": { + "X-Authorization": "$KEY", + }, + "req_template_json_object": { + "text": "$INPUT" + }, + "response_json": "true", + "response_template_json_object": { + "id": "$1$77777$1$" + } + }, + "second_stage": { + "name": "response service", + "uri": "https://example.ai/llm?id=$1$id$1$", + "method": "GET", + "headers": { + "X-Authorization": "1231321", + }, + "response_json": "true", + "response_template_json_object": { + "message": "$OUT$Hello World!$OUT$" + } + } + } + } + } + } @pytest.fixture -def mr_gen(tmp_path): +def mr_gen_burpfile_config(): + _config.plugins.generators["multirest"] = {} + _config.plugins.generators["multirest"]["MultiRestGenerator"] = EXAMPLE_CONFIG["multirest"]["MultiRestGenerator"] generator = MultiRestGenerator() - generator.burpfile = BURP_XML_PATH return generator -def test_extract_output_from_json(): - gen = MultiRestGenerator() - template = {"message": "$OUTPUT"} - real = {"message": "error details"} - outputs = gen.extract_output_from_json(template, real) - assert outputs == ["error details"] + + +@pytest.fixture +def mr_gen_burpfile(): + _config.plugins.generators["multirest"] = {} + _config.plugins.generators["multirest"]["MultiRestGenerator"] = EXAMPLE_BURP_XML + generator = MultiRestGenerator() + return generator -def test_run_extracts_output(mr_gen, requests_mock): - # Prepare the generator with the test XML - mr_gen.reqresp_pairs = mr_gen.get_reqrep_fromburp(BURP_XML_PATH) +# Test run function using burpfile +@pytest.mark.usefixtures("mr_gen_burpfile") +def test_run_extracts_output(mr_gen_burpfile, requests_mock): + mr_gen_burpfile.reqresp_pairs = mr_gen_burpfile.get_reqrep_fromburp() - # Mock the endpoint for both test cases url = "https://example.com/score" - # First response: 200 OK - requests_mock.post(url, [ - { - "status_code": 200, - "json": {"id": "999999"} - }, - ]) + requests_mock.post(url,headers={"Content-Type": "application/json"},status_code= 200,text='{"id": "999999"}') url2 = "https://example.com/score?id=999999" - requests_mock.get(url2, status_code=400, json={"message": "error details"}) + requests_mock.get(url2, headers={"Content-Type": "application/json"},status_code=400, text='{"message": "error details"}') - # Run the generator for the first request (should not extract $OUTPUT) - output = mr_gen.run("Hello") - # The second request in the XML expects $OUTPUT to be extracted + output = mr_gen_burpfile.run("Hello") assert output == "error details" -def test_substitute_placeholders(mr_gen): - mr_gen = MultiRestGenerator() - mr_gen.variables = {"1": "foo"} - text = '{"data":"$1$Hello$1$"}' - result = mr_gen.substitute_placeholders(text) - assert result == '{"data":"foo"}' +# Test run function using the config +@pytest.mark.usefixtures("mr_gen_burpfile_config") +def test_run_with_config_and_mock(mr_gen_burpfile_config, requests_mock): + url = "https://example.ai/llm" + requests_mock.post(url,headers={"Content-Type": "application/json"},status_code= 200,text='{"id": "999999"}') + url2 = "https://example.ai/llm?id=999999" + requests_mock.get(url2, headers={"Content-Type": "application/json"}, status_code=400, json={"message": "error details"}) -def test_response_mismatch_raises(mr_gen, requests_mock): - mr_gen.reqresp_pairs = [{ - "request": { - "method": "GET", - "headers": {"Host": "example.com"}, - "endpoint": "/fail", - "body": "" - }, - "response": { - "body": '{"expected":"value"}', - "headers": {} - } - }] - url = "https://example.com/fail" - requests_mock.get(url, json={"unexpected": "different"}) - with pytest.raises(ValueError): - mr_gen.run("prompt") - -def test_substitute_placeholders_no_vars(mr_gen): - mr_gen.variables = {} - text = '{"data":"$1$Hello$1$"}' - result = mr_gen.substitute_placeholders(text) - assert result == '{"data":"$1$Hello$1$"}' - -def test_extract_vars_from_headers(mr_gen): - packet = { - "headers": {"X-Token": "$1$tokenval$1$"}, - "body": "" - } - vars_found = mr_gen.extract_vars_from_packet(packet) - assert vars_found == {"1": "tokenval"} - -def test_non_json_response(mr_gen, requests_mock): - packet = { - "method": "GET", - "headers": {"Host": "example.com"}, - "endpoint": "/plain", - "body": "" - } - url = "https://example.com/plain" - requests_mock.get(url, text="plain text response") - resp = mr_gen.make_request(packet) - assert resp.text == "plain text response" - -def test_get_request_with_placeholder(mr_gen, requests_mock): - mr_gen.variables = {"1": "abc123"} - packet = { - "method": "GET", - "headers": {"Host": "example.com"}, - "endpoint": "/score?id=$1$foo$1$", - "body": "" - } - url = "https://example.com/score?id=abc123" - requests_mock.get(url, text="ok") - req_packet = mr_gen.substitute_in_packet(packet) - resp = mr_gen.make_request(req_packet) - assert resp.text == "ok" \ No newline at end of file + output = mr_gen_burpfile_config.run("Hello") + assert output == "error details" \ No newline at end of file From 0282df71a504231b69a4e940cd77ae372c6a1954 Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Wed, 23 Jul 2025 20:40:02 -0500 Subject: [PATCH 06/12] Added classes to rest package --- garak/generators/multirest.py | 349 ------------------------ garak/generators/rest.py | 418 ++++++++++++++++++++++++++--- garak/generators/wsocket.py | 47 ---- tests/generators/test_multirest.py | 79 +++--- tests/generators/test_websocket.py | 26 +- tools/rest/multirest.xml | 104 ------- 6 files changed, 437 insertions(+), 586 deletions(-) delete mode 100644 garak/generators/multirest.py delete mode 100644 garak/generators/wsocket.py delete mode 100644 tools/rest/multirest.xml diff --git a/garak/generators/multirest.py b/garak/generators/multirest.py deleted file mode 100644 index 979f7ef75..000000000 --- a/garak/generators/multirest.py +++ /dev/null @@ -1,349 +0,0 @@ -import json -import re - -from typing import List, Union -from garak.generators.base import Generator -from garak import _config - -import xml.etree.ElementTree as ET - -import requests - - -DEFAULT_CLASS = "MultiRestGenerator" - - -class MultiRestGenerator(Generator): - """ - This is a generator to work with multiple request-response pairs using Burp Suite's exported XML format. It supports variable substitution, in case you need to build succeeding requests uponinformation from previous request/response pairs. - The generator supports only JSON request and responses. - - To set a variable in the response, use the format $n$...$n$, where n is a number. - - When you save a Burpfile with reqeust and response items, be sure to uncheck the base64 button. - - For example: - HTTP/2 200 OK - Date: Tue, 20 May 2025 05:58:32 GMT - Content-Type: application/json - Content-Length: 52 - Server: uvicorn - Apigw-Requestid: $2$K2kBRiewPHcEM8A=$2$ - Access-Control-Allow-Origin: https://platform.dreadnode.io - - {"id":"$1$999999$1$"}]]> - - The above response will set the variable 1 to "999999" and variable 2 to "K2kBRiewPHcEM8A=" - - From there you can place these variables in the request body or headers like this: - GET /score?id=$1$123$1$ HTTP/2 - Host: example.com - User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:138.0) Gecko/20100101 Firefox/138.0 - Accept: */* - Accept-Language: en-US,en;q=0.5 - Accept-Encoding: gzip, deflate, br - Referer: https://example.com/score - Content-Type: application/json - Origin: https://example.com - Sec-Fetch-Dest: empty - Sec-Fetch-Mode: cors - Sec-Fetch-Site: same-site - Dnt: 1 - Sec-Gpc: 1 - Priority: u=0 - Te: trailers - - The above request will replace $1$123$1$ with the value of variable 1, which is "999999". - - You can specify where the $INPUT and $OUTPUT variables go in the request and response, as well: - - POST /score HTTP/2 - Host: example.com - User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:138.0) Gecko/20100101 Firefox/138.0 - Accept: */* - Accept-Language: en-US,en;q=0.5 - Accept-Encoding: gzip, deflate, br - Content-Type: application/json - Sec-Fetch-Mode: cors - Sec-Fetch-Site: same-site - Dnt: 1 - Sec-Gpc: 1 - Priority: u=0 - Te: trailers - - {"data":"$INPUT"} - - Here $INPUT will be replaced with the input value provided to call_model. - - HTTP/2 400 Bad Request - Date: Tue, 20 May 2025 05:58:14 GMT - Content-Type: application/json - Content-Length: 46 - Server: uvicorn - Apigw-Requestid: K2j-dj6KPHcEMvw= - Access-Control-Allow-Origin: https://platform.dreadnode.io - - {"message":"$OUTPUT"}]]> - - Here the location of $OUTPUT in the response body indicates that the output value will be extracted from the response JSON under the "message" key. - If the response is not JSON, it will be returned as is. - - You can also specify request-response pairs in a config file similar to that of the Rest Generator: - { - "multi_endpoint_rest": { - "MultiEndpointGenerator": { - "first_stage": { - "name": "request service", - "uri": "https://example.ai/llm", - "method": "post", - "headers": { - "X-Authorization": "$KEY", - }, - "req_template_json_object": { - "text": "$INPUT" - }, - "response_json": true, - "response_json_body": "$1$Hello$1$" - }, - "second_stage": { - "name": "response service", - "uri": "https://example.ai/llm", - "method": "post", - "headers": { - "X-Authorization": "$KEY", - }, - "req_template_json_object": { - "text": "$INPUT" - }, - "response_json": true, - "response_template_json_object": { - "text": "$OUT$data$OUT$" - } - } - } - } - } - """ - - DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { - "reqresp_pairs": [], - "burpfile": None, - "manual": [], - } - - def __init__(self, name="MultiRestGenerator", config_root=_config): - super().__init__(name, config_root=config_root) - # Error condition for when neither burpfile nor manual is provided - if not self.burpfile and not self.manual: - raise ValueError( - "You must provide either a burpfile or manually define request and responses in the config file for MultiRestGenerator." - ) - # Error condition for when both burpfile and manual is provided - if self.burpfile and self.manual: - raise ValueError( - "You cannot provide both a burpfile and manual definitions to the MultiRestGenerator." - ) - # Load request-response pairs from burpfile or manual - if self.burpfile: - self.reqresp_pairs = self.get_reqrep_fromburp() - else: - self.reqresp_pairs = self.get_reqrep_fromconfig() - - def get_reqrep_fromconfig(self): - def parse_domain(uri): - # Extracts the domain from a URI - match = re.match(r"https?://([^/]+)", uri) - return match.group(1) if match else None - - # Parses the manual definitions and returns a list of request-response pairs - pairs = [] - for k, item in self.manual.items(): - # Might need to just move this to a separate function. - req = { - "headers": item["headers"] | {"Host": parse_domain(item["uri"])}, - "body": ( - item["req_template_json_object"] if item["method"] == "post" else {} - ), - "method": item["method"], - "endpoint": item["uri"][item["uri"].rindex("/") :], - } - resp = { - "headers": {}, - "body": ( - item["response_template_json_object"] - if item.get("response_json") - else {} - ), - "status": "", - "error message": "", - } - pairs.append({"request": req, "response": resp}) - return pairs - - def get_reqrep_fromburp(self): - # Parses the Burp Suite XML file and extracts request-response pairs - tree = ET.parse(self.burpfile) - root = tree.getroot() - pairs = [] - for item in root.findall("item"): - req = item.find("request").text - resp = item.find("response").text - pairs.append( - { - "request": self.make_reqrep_dictionary(req), - "response": self.make_reqrep_dictionary(resp), - } - ) - return pairs - - def make_reqrep_dictionary(self, text: str): - # Converts a raw HTTP request or response text into a dictionary format - packet = {"headers": {}, "body": ""} - http_line = text.split("\n") - border = http_line.index("") - for substring in http_line[:border]: - # This condition assumes that the line is a Header - if ":" in substring: - i = substring.index(":") - packet["headers"][substring[:i]] = ( - substring[i + 1 :].rstrip("\n").lstrip(" ").lower() - ) - # This condition assumes that the line is a Request - elif " HTTP/" in substring: - a = substring.rstrip("\n").split(" ") - packet["method"] = a[0] - packet["endpoint"] = a[1] - # This condition assumes that the line is a Response - elif "HTTP/" in substring: - a = substring.rstrip("\n").split(" ") - packet["status"] = a[1] - packet["error message"] = "".join(a[2:]) - try: - packet["body"] = json.loads(http_line[-1]) - except json.decoder.JSONDecodeError: - packet["body"] = http_line[-1] - return packet - - def make_request(self, packet): - method = packet.get("method", "POST") - url = f"https://{packet['headers'].get('Host')}{packet.get('endpoint', '')}" - headers = { - k: v - for k, v in packet["headers"].items() - if k not in ["method", "Host", "endpoint"] - } - data = packet.get("body", "") - if method == "GET": - resp = requests.get(url, headers=headers) - else: - try: - json_data = json.loads(data) - resp = requests.post(url, headers=headers, json=json_data) - except Exception: - resp = requests.post(url, headers=headers, data=data) - return self.parse_response(resp) - - def compare_responses(self, expected, actual): - # This just checks if the keys match. - e_keys = list(expected.keys()).sort() - a_keys = list(actual.keys()).sort() - return e_keys == a_keys - - def extract_variable_locations(self, expected, location=None, locations=None): - if location == None: - location = [] - if locations == None: - locations = {} - - if type(expected) == dict: - for k, v in expected.items(): - self.extract_variable_locations(v, location + [k], locations) - else: - placeholders = re.findall(r"\$(OUT|\d+)\$", expected) - for placeholder in placeholders: - locations[placeholder] = location - return locations - - def update_locations_dictionary(self, dictionary, packet): - location_dict = dictionary.copy() - result = packet - for var_number, locations in location_dict.items(): - if type(locations) == list: - for location in locations: - result = result[location] - else: - continue - location_dict[var_number] = result - return location_dict - - def place_var_into_request_packet(self, packet, locations, new_packet=None): - if not new_packet: - new_packet = {} - if type(packet) == dict: - for k, v in packet.items(): - new_packet[k] = self.place_var_into_request_packet(v, locations) - return new_packet - else: - string_to_edit = packet # At this point, this should be a string. - - def repl(match): - num = match.group(1) - return locations.get(num, match.group(0)) - - return re.sub(r"\$(\d+)\$.*?\$\1\$", repl, string_to_edit) - - def place_input_into_request_packet( - self, packet: dict, prompt: str, new_packet=None - ): - if not new_packet: - new_packet = {} - if type(packet) == dict: - for k, v in packet.items(): - new_packet[k] = self.place_input_into_request_packet(v, prompt) - return new_packet - else: - string_to_compare = packet # At this point, this should be a string. - if string_to_compare == "$INPUT": - return prompt - else: - return string_to_compare - - def parse_response(self, raw_response: requests.Response) -> dict: - return { - "headers": raw_response.headers, - "body": ( - raw_response.json() - if raw_response.headers["Content-Type"] == "application/json" - else raw_response.text - ), - } - - def run(self, prompt: str,locations = {}) -> str: - - for pair in self.reqresp_pairs: - req_packet = self.place_input_into_request_packet(pair["request"], prompt) - req_packet = self.place_var_into_request_packet(req_packet, locations) - resp_packet = self.make_request(req_packet) - if self.compare_responses(pair["response"], resp_packet): - try: - locations.update(self.extract_variable_locations(pair["response"])) - locations.update( - self.update_locations_dictionary(locations, resp_packet) - ) - except: - raise KeyError( - f"Locator failure: Variable could not be found in actual response: {locations} -> {resp_packet}" - ) - else: - raise ValueError( - f"Response mismatch: expected {pair["response"]}, got {resp_packet}" - ) - if locations.get("OUT"): - return locations["OUT"] - else: - return "" - - def _call_model( - self, prompt: str, generations_this_call: int = 1 - ) -> List[Union[str, None]]: - return ["".join(self.run(prompt))] diff --git a/garak/generators/rest.py b/garak/generators/rest.py index e9f4d9005..16e39bbe0 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -6,17 +6,25 @@ Generic Module for REST API connections """ + import json +import xml.etree.ElementTree as ET import logging from typing import List, Union + import requests +import websockets +from websockets.sync.client import connect +import re import backoff import jsonpath_ng from jsonpath_ng.exceptions import JsonPathParserError +# from garak.generators.multirest import MultiRestGenerator + from garak import _config -from garak.exception import APIKeyMissingError, BadGeneratorException, RateLimitHit, GarakBackoffTrigger +from garak.exception import APIKeyMissingError, RateLimitHit from garak.generators.base import Generator @@ -35,8 +43,6 @@ class RestGenerator(Generator): "response_json_field": None, "req_template": "$INPUT", "request_timeout": 20, - "proxies": None, - "verify_ssl": True, } ENV_VAR = "REST_API_KEY" @@ -59,17 +65,14 @@ class RestGenerator(Generator): "request_timeout", "ratelimit_codes", "skip_codes", - "skip_seq_start", - "skip_seq_end", "temperature", "top_k", - "proxies", - "verify_ssl", ) def __init__(self, uri=None, config_root=_config): self.uri = uri self.name = uri + self.seed = _config.run.seed self.supports_multiple_generations = False # not implemented yet self.escape_function = self._json_escape self.retry_5xx = True @@ -123,18 +126,6 @@ def __init__(self, uri=None, config_root=_config): self.method = "post" self.http_function = getattr(requests, self.method) - # validate proxies formatting - # sanity check only leave actual parsing of values to the `requests` library on call. - if hasattr(self, "proxies") and self.proxies is not None: - if not isinstance(self.proxies, dict): - raise BadGeneratorException( - "`proxies` value provided is not in the required format. See documentation from the `requests` package for details on expected format. https://requests.readthedocs.io/en/latest/user/advanced/#proxies" - ) - - # suppress warnings about intentional SSL validation suppression - if isinstance(self.verify_ssl, bool) and not self.verify_ssl: - requests.packages.urllib3.disable_warnings() - # validate jsonpath if self.response_json and self.response_json_field: try: @@ -185,7 +176,8 @@ def _populate_template( output = output.replace("$KEY", self.api_key) return output.replace("$INPUT", self.escape_function(text)) - @backoff.on_exception(backoff.fibo, (RateLimitHit, GarakBackoffTrigger), max_value=70) + # we'll overload IOError as the rate limit exception + @backoff.on_exception(backoff.fibo, RateLimitHit, max_value=70) def _call_model( self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: @@ -209,19 +201,8 @@ def _call_model( data_kw: request_data, "headers": request_headers, "timeout": self.request_timeout, - "proxies": self.proxies, - "verify": self.verify_ssl, } - try: - resp = self.http_function(self.uri, **req_kArgs) - except UnicodeEncodeError as uee: - # only RFC2616 (latin-1) is guaranteed - # don't print a repr, this might leak api keys - logging.error( - "Only latin-1 encoding supported by HTTP RFC 2616, check headers and values for unusual chars", - exc_info=uee, - ) - raise BadGeneratorException from uee + resp = self.http_function(self.uri, **req_kArgs) if resp.status_code in self.skip_codes: logging.debug( @@ -250,7 +231,7 @@ def _call_model( if str(resp.status_code)[0] == "5": error_msg = f"REST URI server error: {resp.status_code} - {resp.reason}, uri: {self.uri}" if self.retry_5xx: - raise GarakBackoffTrigger(error_msg) + raise IOError(error_msg) raise ConnectionError(error_msg) if not self.response_json: @@ -297,3 +278,374 @@ def _call_model( DEFAULT_CLASS = "RestGenerator" + + +class MultiRestGenerator(Generator): + """ + This is a generator to work with multiple request-response pairs using Burp Suite's exported XML format. It supports variable substitution, in case you need to build succeeding requests uponinformation from previous request/response pairs. + The generator supports only JSON request and responses. + + To set a variable in the response, use the format $n$...$n$, where n is a number. + + When you save a Burpfile with reqeust and response items, be sure to uncheck the base64 button. + + For example: + HTTP/2 200 OK + Date: Tue, 20 May 2025 05:58:32 GMT + Content-Type: application/json + Content-Length: 52 + Server: uvicorn + Apigw-Requestid: $2$K2kBRiewPHcEM8A=$2$ + Access-Control-Allow-Origin: https://platform.dreadnode.io + + {"id":"$1$999999$1$"}]]> + + The above response will set the variable 1 to "999999" and variable 2 to "K2kBRiewPHcEM8A=" + + From there you can place these variables in the request body or headers like this: + GET /score?id=$1$123$1$ HTTP/2 + Host: example.com + User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:138.0) Gecko/20100101 Firefox/138.0 + Accept: */* + Accept-Language: en-US,en;q=0.5 + Accept-Encoding: gzip, deflate, br + Referer: https://example.com/score + Content-Type: application/json + Origin: https://example.com + Sec-Fetch-Dest: empty + Sec-Fetch-Mode: cors + Sec-Fetch-Site: same-site + Dnt: 1 + Sec-Gpc: 1 + Priority: u=0 + Te: trailers + + The above request will replace $1$123$1$ with the value of variable 1, which is "999999". + + You can specify where the $INPUT and $OUTPUT variables go in the request and response, as well: + + POST /score HTTP/2 + Host: example.com + User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:138.0) Gecko/20100101 Firefox/138.0 + Accept: */* + Accept-Language: en-US,en;q=0.5 + Accept-Encoding: gzip, deflate, br + Content-Type: application/json + Sec-Fetch-Mode: cors + Sec-Fetch-Site: same-site + Dnt: 1 + Sec-Gpc: 1 + Priority: u=0 + Te: trailers + + {"data":"$INPUT"} + + Here $INPUT will be replaced with the input value provided to call_model. + + HTTP/2 400 Bad Request + Date: Tue, 20 May 2025 05:58:14 GMT + Content-Type: application/json + Content-Length: 46 + Server: uvicorn + Apigw-Requestid: K2j-dj6KPHcEMvw= + Access-Control-Allow-Origin: https://platform.dreadnode.io + + {"message":"$OUTPUT"}]]> + + Here the location of $OUTPUT in the response body indicates that the output value will be extracted from the response JSON under the "message" key. + If the response is not JSON, it will be returned as is. + + You can also specify request-response pairs in a config file similar to that of the Rest Generator: + { + "multi_endpoint_rest": { + "MultiEndpointGenerator": { + "first_stage": { + "name": "request service", + "uri": "https://example.ai/llm", + "method": "post", + "headers": { + "X-Authorization": "$KEY", + }, + "req_template_json_object": { + "text": "$INPUT" + }, + "response_json": true, + "response_json_body": "$1$Hello$1$" + }, + "second_stage": { + "name": "response service", + "uri": "https://example.ai/llm", + "method": "post", + "headers": { + "X-Authorization": "$KEY", + }, + "req_template_json_object": { + "text": "$INPUT" + }, + "response_json": true, + "response_template_json_object": { + "text": "$OUT$data$OUT$" + } + } + } + } + } + """ + + DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { + "reqresp_pairs": [], + "burpfile": None, + "manual": [], + } + + def __init__(self, name="MultiRestGenerator", config_root=_config): + super().__init__(name, config_root=config_root) + # Error condition for when neither burpfile nor manual is provided + if not self.burpfile and not self.manual: + raise ValueError( + "You must provide either a burpfile or manually define request and responses in the config file for MultiRestGenerator." + ) + # Error condition for when both burpfile and manual is provided + if self.burpfile and self.manual: + raise ValueError( + "You cannot provide both a burpfile and manual definitions to the MultiRestGenerator." + ) + # Load request-response pairs from burpfile or manual + if self.burpfile: + self.reqresp_pairs = self.get_reqrep_fromburp() + else: + self.reqresp_pairs = self.get_reqrep_fromconfig() + + def get_reqrep_fromconfig(self): + def parse_domain(uri): + # Extracts the domain from a URI + match = re.match(r"https?://([^/]+)", uri) + return match.group(1) if match else None + + # Parses the manual definitions and returns a list of request-response pairs + pairs = [] + for k, item in self.manual.items(): + # Might need to just move this to a separate function. + req = { + "headers": item["headers"] | {"Host": parse_domain(item["uri"])}, + "body": ( + item["req_template_json_object"] if item["method"] == "post" else {} + ), + "method": item["method"], + "endpoint": item["uri"][item["uri"].rindex("/") :], + } + resp = { + "headers": {}, + "body": ( + item["response_template_json_object"] + if item.get("response_json") + else {} + ), + "status": "", + "error message": "", + } + pairs.append({"request": req, "response": resp}) + return pairs + + def get_reqrep_fromburp(self): + # Parses the Burp Suite XML file and extracts request-response pairs + tree = ET.parse(self.burpfile) + root = tree.getroot() + pairs = [] + for item in root.findall("item"): + req = item.find("request").text + resp = item.find("response").text + pairs.append( + { + "request": self.make_reqrep_dictionary(req), + "response": self.make_reqrep_dictionary(resp), + } + ) + return pairs + + def make_reqrep_dictionary(self, text: str): + # Converts a raw HTTP request or response text into a dictionary format + packet = {"headers": {}, "body": ""} + http_line = text.split("\n") + border = http_line.index("") + for substring in http_line[:border]: + # This condition assumes that the line is a Header + if ":" in substring: + i = substring.index(":") + packet["headers"][substring[:i]] = ( + substring[i + 1 :].rstrip("\n").lstrip(" ").lower() + ) + # This condition assumes that the line is a Request + elif " HTTP/" in substring: + a = substring.rstrip("\n").split(" ") + packet["method"] = a[0] + packet["endpoint"] = a[1] + # This condition assumes that the line is a Response + elif "HTTP/" in substring: + a = substring.rstrip("\n").split(" ") + packet["status"] = a[1] + packet["error message"] = "".join(a[2:]) + try: + packet["body"] = json.loads(http_line[-1]) + except json.decoder.JSONDecodeError: + packet["body"] = http_line[-1] + return packet + + def make_request(self, packet): + method = packet.get("method", "POST") + url = f"https://{packet['headers'].get('Host')}{packet.get('endpoint', '')}" + headers = { + k: v + for k, v in packet["headers"].items() + if k not in ["method", "Host", "endpoint"] + } + data = packet.get("body", "") + if method == "GET": + resp = requests.get(url, headers=headers) + else: + try: + json_data = json.loads(data) + resp = requests.post(url, headers=headers, json=json_data) + except Exception: + resp = requests.post(url, headers=headers, data=data) + return self.parse_response(resp) + + def compare_responses(self, expected, actual): + # This just checks if the keys match. + e_keys = list(expected.keys()).sort() + a_keys = list(actual.keys()).sort() + return e_keys == a_keys + + def extract_variable_locations(self, expected, location=None, locations=None): + if location == None: + location = [] + if locations == None: + locations = {} + + if type(expected) == dict: + for k, v in expected.items(): + self.extract_variable_locations(v, location + [k], locations) + else: + placeholders = re.findall(r"\$(OUT|\d+)\$", expected) + for placeholder in placeholders: + locations[placeholder] = location + return locations + + def update_locations_dictionary(self, dictionary, packet): + location_dict = dictionary.copy() + result = packet + for var_number, locations in location_dict.items(): + if type(locations) == list: + for location in locations: + result = result[location] + else: + continue + location_dict[var_number] = result + return location_dict + + def place_var_into_request_packet(self, packet, locations, new_packet=None): + if not new_packet: + new_packet = {} + if type(packet) == dict: + for k, v in packet.items(): + new_packet[k] = self.place_var_into_request_packet(v, locations) + return new_packet + else: + string_to_edit = packet # At this point, this should be a string. + + def repl(match): + num = match.group(1) + return locations.get(num, match.group(0)) + + return re.sub(r"\$(\d+)\$.*?\$\1\$", repl, string_to_edit) + + def place_input_into_request_packet( + self, packet: dict, prompt: str, new_packet=None + ): + if not new_packet: + new_packet = {} + if type(packet) == dict: + for k, v in packet.items(): + new_packet[k] = self.place_input_into_request_packet(v, prompt) + return new_packet + else: + string_to_compare = packet # At this point, this should be a string. + if string_to_compare == "$INPUT": + return prompt + else: + return string_to_compare + + def parse_response(self, raw_response: requests.Response) -> dict: + return { + "headers": raw_response.headers, + "body": ( + raw_response.json() + if raw_response.headers["Content-Type"] == "application/json" + else raw_response.text + ), + } + + def run(self, prompt: str, locations={}) -> str: + + for pair in self.reqresp_pairs: + req_packet = self.place_input_into_request_packet(pair["request"], prompt) + req_packet = self.place_var_into_request_packet(req_packet, locations) + resp_packet = self.make_request(req_packet) + if self.compare_responses(pair["response"], resp_packet): + try: + locations.update(self.extract_variable_locations(pair["response"])) + locations.update( + self.update_locations_dictionary(locations, resp_packet) + ) + except: + raise KeyError( + f"Locator failure: Variable could not be found in actual response: {locations} -> {resp_packet}" + ) + else: + raise ValueError( + f"Response mismatch: expected {pair["response"]}, got {resp_packet}" + ) + if locations.get("OUT"): + return locations["OUT"] + else: + return "" + + def _call_model( + self, prompt: str, generations_this_call: int = 1 + ) -> List[Union[str, None]]: + return ["".join(self.run(prompt))] + + +class WebSocketGenerator(Generator): + """ + This is a generator to work with websockets + """ + + DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { + "uri": None, + "auth_key": None, + "body": "{}", + "json_response": True, + "json_key": "output", + } + + def __init__(self, name="WebSocket", config_root=_config): + super().__init__(name, config_root=config_root) + + def json_handler(self, data): + response_json = json.loads(data) + return json.dumps(response_json[self.json_key]) + + def request(self, payload): + with connect(self.uri) as websocket: + websocket.send(self.body.replace("$INPUT", payload)) + message = websocket.recv() + return self.json_handler(message) if self.json_response == True else message + + def _call_model( + self, prompt: str, generations_this_call: int = 1 + ) -> List[Union[str, None]]: + if output := self.request(self, prompt) == dict: + return output[self.json_key] + else: + return output diff --git a/garak/generators/wsocket.py b/garak/generators/wsocket.py deleted file mode 100644 index 1011281aa..000000000 --- a/garak/generators/wsocket.py +++ /dev/null @@ -1,47 +0,0 @@ -import json -import logging -import re - -from typing import List, Union -from garak.generators.base import Generator -from garak import _config - -import websockets -from websockets.sync.client import connect - -class WebSocketGenerator(Generator): - """ - This is a generator to work with websockets - """ - - DEFAULT_PARAMS = Generator.DEFAULT_PARAMS | { - "uri": None, - "auth_key": None, - "body": '{}', - "json_response": True, - "json_key": "output", - } - - def __init__(self, name="WebSocket", config_root=_config): - super().__init__(name, config_root=config_root) - - def json_handler(self, data): - response_json = json.loads(data) - return json.dumps(response_json[self.json_key]) - - - def request(self, payload): - with connect(self.uri) as websocket: - websocket.send(self.body.replace("$INPUT", payload)) - message = websocket.recv() - return self.json_handler(message) if self.json_response == True else message - - - def _call_model(self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: - if output := self.request(self, prompt) == dict : - return output[self.json_key] - else : - return output - - -DEFAULT_CLASS = "WebSocketGenerator" \ No newline at end of file diff --git a/tests/generators/test_multirest.py b/tests/generators/test_multirest.py index a68b3a7dd..bc9a5a0a5 100644 --- a/tests/generators/test_multirest.py +++ b/tests/generators/test_multirest.py @@ -1,29 +1,23 @@ import pytest -from garak.generators.multirest import MultiRestGenerator +from garak.generators.rest import MultiRestGenerator from garak import _config, _plugins -EXAMPLE_BURP_XML = { - "burpfile":"./garak/data/rest/multirest.xml" -} +EXAMPLE_BURP_XML = {"burpfile": "./garak/data/rest/multirest.xml"} EXAMPLE_CONFIG = { - "multirest": { - "MultiRestGenerator": { - "manual": { - "first_stage": { - "name": "request service", - "uri": "https://example.ai/llm", - "method": "POST", - "headers": { + "rest": { + "MultiRestGenerator": { + "manual": { + "first_stage": { + "name": "request service", + "uri": "https://example.ai/llm", + "method": "POST", + "headers": { "X-Authorization": "$KEY", }, - "req_template_json_object": { - "text": "$INPUT" - }, + "req_template_json_object": {"text": "$INPUT"}, "response_json": "true", - "response_template_json_object": { - "id": "$1$77777$1$" - } + "response_template_json_object": {"id": "$1$77777$1$"}, }, "second_stage": { "name": "response service", @@ -35,38 +29,51 @@ "response_json": "true", "response_template_json_object": { "message": "$OUT$Hello World!$OUT$" - } - } + }, + }, } } } - } +} + @pytest.fixture def mr_gen_burpfile_config(): - _config.plugins.generators["multirest"] = {} - _config.plugins.generators["multirest"]["MultiRestGenerator"] = EXAMPLE_CONFIG["multirest"]["MultiRestGenerator"] + _config.plugins.generators["rest"] = {} + _config.plugins.generators["rest"]["MultiRestGenerator"] = EXAMPLE_CONFIG["rest"][ + "MultiRestGenerator" + ] generator = MultiRestGenerator() return generator - @pytest.fixture def mr_gen_burpfile(): - _config.plugins.generators["multirest"] = {} - _config.plugins.generators["multirest"]["MultiRestGenerator"] = EXAMPLE_BURP_XML + _config.plugins.generators["rest"] = {} + _config.plugins.generators["rest"]["MultiRestGenerator"] = EXAMPLE_BURP_XML generator = MultiRestGenerator() return generator + # Test run function using burpfile @pytest.mark.usefixtures("mr_gen_burpfile") def test_run_extracts_output(mr_gen_burpfile, requests_mock): mr_gen_burpfile.reqresp_pairs = mr_gen_burpfile.get_reqrep_fromburp() url = "https://example.com/score" - requests_mock.post(url,headers={"Content-Type": "application/json"},status_code= 200,text='{"id": "999999"}') + requests_mock.post( + url, + headers={"Content-Type": "application/json"}, + status_code=200, + text='{"id": "999999"}', + ) url2 = "https://example.com/score?id=999999" - requests_mock.get(url2, headers={"Content-Type": "application/json"},status_code=400, text='{"message": "error details"}') + requests_mock.get( + url2, + headers={"Content-Type": "application/json"}, + status_code=400, + text='{"message": "error details"}', + ) output = mr_gen_burpfile.run("Hello") assert output == "error details" @@ -76,9 +83,19 @@ def test_run_extracts_output(mr_gen_burpfile, requests_mock): @pytest.mark.usefixtures("mr_gen_burpfile_config") def test_run_with_config_and_mock(mr_gen_burpfile_config, requests_mock): url = "https://example.ai/llm" - requests_mock.post(url,headers={"Content-Type": "application/json"},status_code= 200,text='{"id": "999999"}') + requests_mock.post( + url, + headers={"Content-Type": "application/json"}, + status_code=200, + text='{"id": "999999"}', + ) url2 = "https://example.ai/llm?id=999999" - requests_mock.get(url2, headers={"Content-Type": "application/json"}, status_code=400, json={"message": "error details"}) + requests_mock.get( + url2, + headers={"Content-Type": "application/json"}, + status_code=400, + json={"message": "error details"}, + ) output = mr_gen_burpfile_config.run("Hello") - assert output == "error details" \ No newline at end of file + assert output == "error details" diff --git a/tests/generators/test_websocket.py b/tests/generators/test_websocket.py index 1f02249e3..a0f0a2fbe 100644 --- a/tests/generators/test_websocket.py +++ b/tests/generators/test_websocket.py @@ -1,7 +1,8 @@ import pytest from unittest.mock import patch, MagicMock -from garak.generators.wsocket import WebSocketGenerator +from garak.generators.rest import WebSocketGenerator + @pytest.fixture def ws_gen(): @@ -12,34 +13,15 @@ def ws_gen(): gen.json_key = "output" return gen + def test_json_handler(ws_gen): data = '{"output": "test response"}' result = ws_gen.json_handler(data) assert result == '"test response"' -@patch("garak.generators.wsocket.connect") -def test_request_json_response(mock_connect, ws_gen): - mock_ws = MagicMock() - mock_ws.recv.return_value = '{"output": "foo"}' - mock_connect.return_value.__enter__.return_value = mock_ws - - result = ws_gen.request("bar") - assert result == '"foo"' - mock_ws.send.assert_called_once_with('{"input": "bar"}') - -@patch("garak.generators.wsocket.connect") -def test_request_raw_response(mock_connect, ws_gen): - ws_gen.json_response = False - mock_ws = MagicMock() - mock_ws.recv.return_value = "raw" - mock_connect.return_value.__enter__.return_value = mock_ws - - result = ws_gen.request("baz") - assert result == "raw" def test_live_request_raw_response(ws_gen): ws_gen.json_response = False ws_gen.uri = "wss://echo.websocket.events" result = ws_gen.request("test") - assert result == "test" - \ No newline at end of file + assert "echo.websocket.events" in result diff --git a/tools/rest/multirest.xml b/tools/rest/multirest.xml deleted file mode 100644 index 66987fee2..000000000 --- a/tools/rest/multirest.xml +++ /dev/null @@ -1,104 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - -]> - - - - - example.com - 443 - https - - - null - - 200 - 270 - JSON - - - - - - - example.com - 443 - https - - - null - - 400 - 273 - JSON - - - - From ff5929ecc8cfc43a2db266f4647c718be47e251a Mon Sep 17 00:00:00 2001 From: iamnotcj <69222710+iamnotcj@users.noreply.github.com> Date: Wed, 23 Jul 2025 20:51:57 -0500 Subject: [PATCH 07/12] Added the latest version of the RestGenerator --- garak/generators/rest.py | 49 ++++++++++++++++++++++++++++++++-------- 1 file changed, 40 insertions(+), 9 deletions(-) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index 16e39bbe0..4b2dfa56d 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -13,7 +13,6 @@ from typing import List, Union import requests -import websockets from websockets.sync.client import connect import re @@ -21,10 +20,13 @@ import jsonpath_ng from jsonpath_ng.exceptions import JsonPathParserError -# from garak.generators.multirest import MultiRestGenerator - from garak import _config -from garak.exception import APIKeyMissingError, RateLimitHit +from garak.exception import ( + APIKeyMissingError, + BadGeneratorException, + RateLimitHit, + GarakBackoffTrigger, +) from garak.generators.base import Generator @@ -43,6 +45,8 @@ class RestGenerator(Generator): "response_json_field": None, "req_template": "$INPUT", "request_timeout": 20, + "proxies": None, + "verify_ssl": True, } ENV_VAR = "REST_API_KEY" @@ -65,14 +69,17 @@ class RestGenerator(Generator): "request_timeout", "ratelimit_codes", "skip_codes", + "skip_seq_start", + "skip_seq_end", "temperature", "top_k", + "proxies", + "verify_ssl", ) def __init__(self, uri=None, config_root=_config): self.uri = uri self.name = uri - self.seed = _config.run.seed self.supports_multiple_generations = False # not implemented yet self.escape_function = self._json_escape self.retry_5xx = True @@ -126,6 +133,18 @@ def __init__(self, uri=None, config_root=_config): self.method = "post" self.http_function = getattr(requests, self.method) + # validate proxies formatting + # sanity check only leave actual parsing of values to the `requests` library on call. + if hasattr(self, "proxies") and self.proxies is not None: + if not isinstance(self.proxies, dict): + raise BadGeneratorException( + "`proxies` value provided is not in the required format. See documentation from the `requests` package for details on expected format. https://requests.readthedocs.io/en/latest/user/advanced/#proxies" + ) + + # suppress warnings about intentional SSL validation suppression + if isinstance(self.verify_ssl, bool) and not self.verify_ssl: + requests.packages.urllib3.disable_warnings() + # validate jsonpath if self.response_json and self.response_json_field: try: @@ -176,8 +195,9 @@ def _populate_template( output = output.replace("$KEY", self.api_key) return output.replace("$INPUT", self.escape_function(text)) - # we'll overload IOError as the rate limit exception - @backoff.on_exception(backoff.fibo, RateLimitHit, max_value=70) + @backoff.on_exception( + backoff.fibo, (RateLimitHit, GarakBackoffTrigger), max_value=70 + ) def _call_model( self, prompt: str, generations_this_call: int = 1 ) -> List[Union[str, None]]: @@ -201,8 +221,19 @@ def _call_model( data_kw: request_data, "headers": request_headers, "timeout": self.request_timeout, + "proxies": self.proxies, + "verify": self.verify_ssl, } - resp = self.http_function(self.uri, **req_kArgs) + try: + resp = self.http_function(self.uri, **req_kArgs) + except UnicodeEncodeError as uee: + # only RFC2616 (latin-1) is guaranteed + # don't print a repr, this might leak api keys + logging.error( + "Only latin-1 encoding supported by HTTP RFC 2616, check headers and values for unusual chars", + exc_info=uee, + ) + raise BadGeneratorException from uee if resp.status_code in self.skip_codes: logging.debug( @@ -231,7 +262,7 @@ def _call_model( if str(resp.status_code)[0] == "5": error_msg = f"REST URI server error: {resp.status_code} - {resp.reason}, uri: {self.uri}" if self.retry_5xx: - raise IOError(error_msg) + raise GarakBackoffTrigger(error_msg) raise ConnectionError(error_msg) if not self.response_json: From b82b61d9c61213c9d9bc60892b15f93866d7f043 Mon Sep 17 00:00:00 2001 From: CJ Anih <69222710+iamnotcj@users.noreply.github.com> Date: Fri, 25 Jul 2025 09:17:07 -0500 Subject: [PATCH 08/12] Update tests/generators/test_multirest.py Co-authored-by: Jeffrey Martin Signed-off-by: CJ Anih <69222710+iamnotcj@users.noreply.github.com> --- tests/generators/test_multirest.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/generators/test_multirest.py b/tests/generators/test_multirest.py index bc9a5a0a5..f8e10cdc9 100644 --- a/tests/generators/test_multirest.py +++ b/tests/generators/test_multirest.py @@ -1,8 +1,10 @@ import pytest from garak.generators.rest import MultiRestGenerator from garak import _config, _plugins - -EXAMPLE_BURP_XML = {"burpfile": "./garak/data/rest/multirest.xml"} +from pathlib import Path +from garak.data import path as data_path + +EXAMPLE_BURP_XML = {"burpfile": str(Path(data_path) / "rest" / "multirest.xml") } EXAMPLE_CONFIG = { "rest": { From dacfd6d59bfd995d9189a26d17398c87e34852b1 Mon Sep 17 00:00:00 2001 From: CJ Anih <69222710+iamnotcj@users.noreply.github.com> Date: Fri, 25 Jul 2025 09:19:07 -0500 Subject: [PATCH 09/12] Moved DEFAULT_CLASS to bottom of rest.py Signed-off-by: CJ Anih <69222710+iamnotcj@users.noreply.github.com> --- garak/generators/rest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index 4b2dfa56d..dbe2a5143 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -308,7 +308,6 @@ def _call_model( return response -DEFAULT_CLASS = "RestGenerator" class MultiRestGenerator(Generator): @@ -680,3 +679,5 @@ def _call_model( return output[self.json_key] else: return output + +DEFAULT_CLASS = "RestGenerator" From 55e0cf403fbf99e19757aec7669370db3e56031e Mon Sep 17 00:00:00 2001 From: CJ Anih <69222710+iamnotcj@users.noreply.github.com> Date: Fri, 25 Jul 2025 09:20:33 -0500 Subject: [PATCH 10/12] Update garak/generators/rest.py Co-authored-by: Jeffrey Martin Signed-off-by: CJ Anih <69222710+iamnotcj@users.noreply.github.com> --- garak/generators/rest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index dbe2a5143..2023739a0 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -632,8 +632,9 @@ def run(self, prompt: str, locations={}) -> str: f"Locator failure: Variable could not be found in actual response: {locations} -> {resp_packet}" ) else: + response = pair["response"] raise ValueError( - f"Response mismatch: expected {pair["response"]}, got {resp_packet}" + f"Response mismatch: expected {response}, got {resp_packet}" ) if locations.get("OUT"): return locations["OUT"] From ed75c6e79f8ae2a15e32eb4bf95975810015ff22 Mon Sep 17 00:00:00 2001 From: CJ Anih <69222710+iamnotcj@users.noreply.github.com> Date: Fri, 25 Jul 2025 15:06:31 -0500 Subject: [PATCH 11/12] Update tests/generators/test_multirest.py Co-authored-by: Jeffrey Martin Signed-off-by: CJ Anih <69222710+iamnotcj@users.noreply.github.com> --- tests/generators/test_multirest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/generators/test_multirest.py b/tests/generators/test_multirest.py index f8e10cdc9..2b6593e57 100644 --- a/tests/generators/test_multirest.py +++ b/tests/generators/test_multirest.py @@ -4,7 +4,7 @@ from pathlib import Path from garak.data import path as data_path -EXAMPLE_BURP_XML = {"burpfile": str(Path(data_path) / "rest" / "multirest.xml") } +EXAMPLE_BURP_XML = {"burpfile": str(_config.transient.package_dir / "data" / "rest" / "multirest.xml") } EXAMPLE_CONFIG = { "rest": { From a4cd4fb2ebd8655046c0922379baa040d7471003 Mon Sep 17 00:00:00 2001 From: CJ Anih <69222710+iamnotcj@users.noreply.github.com> Date: Fri, 1 Aug 2025 13:53:13 -0500 Subject: [PATCH 12/12] Update documentation string for MultiRestGenerator Signed-off-by: CJ Anih <69222710+iamnotcj@users.noreply.github.com> --- garak/generators/rest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/garak/generators/rest.py b/garak/generators/rest.py index 2023739a0..fbd36c062 100644 --- a/garak/generators/rest.py +++ b/garak/generators/rest.py @@ -380,7 +380,7 @@ class MultiRestGenerator(Generator): Apigw-Requestid: K2j-dj6KPHcEMvw= Access-Control-Allow-Origin: https://platform.dreadnode.io - {"message":"$OUTPUT"}]]> + {"message":"$OUT$Hello World$OUT$"}]]> Here the location of $OUTPUT in the response body indicates that the output value will be extracted from the response JSON under the "message" key. If the response is not JSON, it will be returned as is.