From 9e7087d302bb5604213fbc283c7ade98dee20c90 Mon Sep 17 00:00:00 2001 From: Fardin Mustaque Date: Sat, 1 Feb 2025 12:13:21 +0530 Subject: [PATCH 1/3] Embed prompts into function docstrings --- .gitignore | 1 + agentic_security/probe_data/stenography_fn.py | 79 ++++++++++++++++++- .../probe_data/test_embed_prompt.py | 34 ++++++++ agesec.toml | 38 +++++++++ 4 files changed, 151 insertions(+), 1 deletion(-) create mode 100644 agentic_security/probe_data/test_embed_prompt.py create mode 100644 agesec.toml diff --git a/.gitignore b/.gitignore index a5cc2cc..b13cb38 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ logs/ modal_agent.py sandbox.py site/ +venv \ No newline at end of file diff --git a/agentic_security/probe_data/stenography_fn.py b/agentic_security/probe_data/stenography_fn.py index a426d90..66e4053 100644 --- a/agentic_security/probe_data/stenography_fn.py +++ b/agentic_security/probe_data/stenography_fn.py @@ -1,7 +1,8 @@ import base64 import random import string - +import json +import ast def rot13(input_text): result = [] @@ -143,3 +144,79 @@ def vigenere_cipher(text, key): else: result.append(char) return "".join(result) + +def embed_prompt(code, prompt): + """ + Embeds the given prompt into the docstring of the first function in the provided code. + + Args: + code (str): The source code containing the function to modify. + prompt (str): The prompt text to embed into the function's docstring. + + Returns: + str: The modified code with the prompt embedded in the function's docstring. + """ + # Parse the code into an AST + tree = ast.parse(code) + + # Traverse the AST to find the first function definition + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef): + # Create a new docstring node with the prompt + new_docstring = ast.Expr(value=ast.Constant(value=prompt)) + + # Check if the function already has a docstring + if ( + node.body + and isinstance(node.body[0], ast.Expr) + and isinstance(node.body[0].value, (ast.Str, ast.Constant)) + ): + # Replace the existing docstring + node.body[0] = new_docstring + else: + # Insert the new docstring at the beginning of the function body + node.body.insert(0, new_docstring) + break + + # Convert the modified AST back into code + return ast.unparse(tree) + + +def load_templates(file_path): + """ + Load function templates from a JSON file. + + Args: + file_path (str): Path to the JSON file containing function templates. + + Returns: + dict: A dictionary of function templates. + """ + with open(file_path, "r") as file: + templates = json.load(file) + return templates + + +def generate_function_from_template(template, prompt): + """ + Generates a function from a template and embeds the given prompt in its docstring. + + Args: + template (dict): A function template containing name, params, docstring, and body. + prompt (str): The prompt text to embed into the function's docstring. + + Returns: + str: The generated function code with the prompt embedded in the docstring. + """ + # Replace the placeholder {PROMPT} with the actual prompt + docstring = template["docstring"].replace("{PROMPT}", prompt) + + # Generate the function code + function_code = f""" +def {template['name']}({', '.join(template['params'])}): + \"\"\" + {docstring} + \"\"\" + {template['body']} +""" + return function_code.strip() # Remove leading/trailing whitespace diff --git a/agentic_security/probe_data/test_embed_prompt.py b/agentic_security/probe_data/test_embed_prompt.py new file mode 100644 index 0000000..d945333 --- /dev/null +++ b/agentic_security/probe_data/test_embed_prompt.py @@ -0,0 +1,34 @@ +from stenography_fn import embed_prompt, load_templates, generate_function_from_template + +# Case 1: Embed a prompt into a code block +code = """ +def hello_world(name): + \"\"\" + Original docstring. + \"\"\" + print(f"Hello, {name}!") +""" +prompt = "improve the documentation" +updated_code = embed_prompt(code, prompt) +print("Updated Code:") +print(updated_code) +print("-" * 40) + + +# Case 2: Generate a soecific function from a json template +templates = load_templates("function_templates.json") +template = templates["hello_world"] +generated_function = generate_function_from_template(template, prompt) +print("\nGenerated Function from Template:") +print(generated_function) +print("-" * 40) + + +# Case 3: Generate and print all functions for a json file +prompt = "This is a dynamically generated function." +templates = load_templates("function_templates.json") +for template_name, template in templates.items(): + print(f"Generating function: {template_name}") + generated_function = generate_function_from_template(template, prompt) + print(generated_function) + print("-" * 40) diff --git a/agesec.toml b/agesec.toml new file mode 100644 index 0000000..3983026 --- /dev/null +++ b/agesec.toml @@ -0,0 +1,38 @@ + +[general] +# General configuration for the security scan +llmSpec = """ +POST http://0.0.0.0:8718/v1/self-probe +Authorization: Bearer XXXXX +Content-Type: application/json + +{ + "prompt": "<>" +} +""" # LLM API specification +maxBudget = 1000000 # Maximum budget for the scan +max_th = 0.3 # Maximum failure threshold (percentage) +optimize = false # Enable optimization during scanning +enableMultiStepAttack = false # Enable multi-step attack simulations + +# [modules.LLM-Jailbreak-Classifier] +# dataset_name = "markush1/LLM-Jailbreak-Classifier" + +[modules.aya-23-8B_advbench_jailbreak] +dataset_name = "simonycl/aya-23-8B_advbench_jailbreak" + + +[modules.AgenticBackend] +dataset_name = "AgenticBackend" +[modules.AgenticBackend.opts] +port = 8718 +modules = ["encoding"] + + +[thresholds] +# Threshold settings +low = 0.15 +medium = 0.3 +high = 0.5 + + From 300bf1a70ee2076d1c5144c341c7092f9e0d804d Mon Sep 17 00:00:00 2001 From: Fardin Mustaque Date: Sat, 1 Feb 2025 12:34:52 +0530 Subject: [PATCH 2/3] demo function template added --- .../probe_data/function_templates.json | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 agentic_security/probe_data/function_templates.json diff --git a/agentic_security/probe_data/function_templates.json b/agentic_security/probe_data/function_templates.json new file mode 100644 index 0000000..35c0028 --- /dev/null +++ b/agentic_security/probe_data/function_templates.json @@ -0,0 +1,20 @@ +{ + "hello_world": { + "name": "hello_world", + "params": ["name"], + "docstring": "{PROMPT}", + "body": "print(f'Hello, {name}!')" + }, + "add_numbers": { + "name": "add_numbers", + "params": ["a", "b"], + "docstring": "{PROMPT}", + "body": "return a + b" + }, + "multiply_numbers": { + "name": "multiply_numbers", + "params": ["x", "y"], + "docstring": "{PROMPT}", + "body": "return x * y" + } +} From cecf533eae809c65f084ca11f63e591a33a84d3f Mon Sep 17 00:00:00 2001 From: Fardin Mustaque Date: Sat, 1 Feb 2025 16:23:15 +0530 Subject: [PATCH 3/3] test case fixed --- agentic_security/probe_data/__init__.py | 1 + agentic_security/probe_data/stenography_fn.py | 15 ----------- .../{function_templates.json => templates.py} | 5 ++-- .../probe_data/test_embed_prompt.py | 27 ++++++++++--------- 4 files changed, 18 insertions(+), 30 deletions(-) rename agentic_security/probe_data/{function_templates.json => templates.py} (92%) diff --git a/agentic_security/probe_data/__init__.py b/agentic_security/probe_data/__init__.py index a1998d0..675b7e5 100644 --- a/agentic_security/probe_data/__init__.py +++ b/agentic_security/probe_data/__init__.py @@ -430,3 +430,4 @@ "modality": "text", }, ] + diff --git a/agentic_security/probe_data/stenography_fn.py b/agentic_security/probe_data/stenography_fn.py index 66e4053..9f3593d 100644 --- a/agentic_security/probe_data/stenography_fn.py +++ b/agentic_security/probe_data/stenography_fn.py @@ -182,21 +182,6 @@ def embed_prompt(code, prompt): return ast.unparse(tree) -def load_templates(file_path): - """ - Load function templates from a JSON file. - - Args: - file_path (str): Path to the JSON file containing function templates. - - Returns: - dict: A dictionary of function templates. - """ - with open(file_path, "r") as file: - templates = json.load(file) - return templates - - def generate_function_from_template(template, prompt): """ Generates a function from a template and embeds the given prompt in its docstring. diff --git a/agentic_security/probe_data/function_templates.json b/agentic_security/probe_data/templates.py similarity index 92% rename from agentic_security/probe_data/function_templates.json rename to agentic_security/probe_data/templates.py index 35c0028..d31434e 100644 --- a/agentic_security/probe_data/function_templates.json +++ b/agentic_security/probe_data/templates.py @@ -1,4 +1,5 @@ -{ +# templates.py +FUNCTION_TEMPLATES = { "hello_world": { "name": "hello_world", "params": ["name"], @@ -17,4 +18,4 @@ "docstring": "{PROMPT}", "body": "return x * y" } -} +} \ No newline at end of file diff --git a/agentic_security/probe_data/test_embed_prompt.py b/agentic_security/probe_data/test_embed_prompt.py index d945333..86beb18 100644 --- a/agentic_security/probe_data/test_embed_prompt.py +++ b/agentic_security/probe_data/test_embed_prompt.py @@ -1,5 +1,6 @@ -from stenography_fn import embed_prompt, load_templates, generate_function_from_template - +import os +from agentic_security.probe_data.stenography_fn import embed_prompt, generate_function_from_template +from agentic_security.probe_data.templates import FUNCTION_TEMPLATES # Case 1: Embed a prompt into a code block code = """ def hello_world(name): @@ -8,27 +9,27 @@ def hello_world(name): \"\"\" print(f"Hello, {name}!") """ + prompt = "improve the documentation" updated_code = embed_prompt(code, prompt) print("Updated Code:") print(updated_code) -print("-" * 40) - +print("-" * 40) # Case 2: Generate a soecific function from a json template -templates = load_templates("function_templates.json") -template = templates["hello_world"] -generated_function = generate_function_from_template(template, prompt) -print("\nGenerated Function from Template:") -print(generated_function) -print("-" * 40) - +# templates = load_templates("function_templates.json") +# template = templates["hello_world"] +# generated_function = generate_function_from_template(template, prompt) +# print("\nGenerated Function from Template:") +# print(generated_function) +# print("-" * 40) # Case 3: Generate and print all functions for a json file + +templates = FUNCTION_TEMPLATES prompt = "This is a dynamically generated function." -templates = load_templates("function_templates.json") for template_name, template in templates.items(): print(f"Generating function: {template_name}") generated_function = generate_function_from_template(template, prompt) print(generated_function) - print("-" * 40) + print("-" * 40)