Skip to content

Commit dccaa84

Browse files
knewbury01knewbury01
authored
Improve agents sdk modelling (#5)
* Add testcase and coverage for agents sdk runner run with input param * Rename agent sdk module for clarity * Add case for unnamed param use in runner run from agent sdk
1 parent 0a36be1 commit dccaa84

File tree

5 files changed

+66
-3
lines changed

5 files changed

+66
-3
lines changed

‎python/ql/lib/semmle/python/frameworks/OpenAI.qll‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,32 @@
11
/**
2-
* Provides classes modeling security-relevant aspects of the `openAI`Agents SDK package.
2+
* Provides classes modeling security-relevant aspects of the `openAI` Agents SDK package.
33
* See https://github.com/openai/openai-agents-python.
4+
* As well as the regular openai python interface.
5+
* See https://github.com/openai/openai-python.
46
*/
57

68
private import python
79
private import semmle.python.ApiGraphs
810

11+
/**
12+
* Provides models for agents SDK (instances of the `agents.Runner` class etc).
13+
*
14+
* See https://github.com/openai/openai-agents-python.
15+
*/
16+
module AgentSDK {
17+
/** Gets a reference to the `agents.Agent` class. */
18+
API::Node classRef() { result = API::moduleImport("agents").getMember("Runner") }
19+
20+
API::Node runMembers() { result = classRef().getMember(["run", "run_sync", "run_streamed"]) }
21+
22+
/** Gets a reference to a potential property of `agents.Runner` called input which can refer to a system prompt depending on the role specified. */
23+
API::Node getContentNode() {
24+
result = runMembers().getKeywordParameter("input").getASubscript().getSubscript("content")
25+
or
26+
result = runMembers().getParameter(_).getASubscript().getSubscript("content")
27+
}
28+
}
29+
930
/**
1031
* Provides models for Agent (instances of the `openai.OpenAI` class).
1132
*

‎python/ql/lib/semmle/python/security/dataflow/PromptInjectionCustomizations.qll‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ module PromptInjection {
5252
private class PromptContentSink extends Sink {
5353
PromptContentSink() {
5454
this = OpenAI::getContentNode().asSink()
55+
or
56+
this = AgentSDK::getContentNode().asSink()
5557
}
5658
}
5759

‎python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.expected‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#select
22
| agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | This prompt construction depends on a $@. | agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | user-provided value |
3+
| agent_instructions.py:25:28:25:32 | ControlFlowNode for input | agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | agent_instructions.py:25:28:25:32 | ControlFlowNode for input | This prompt construction depends on a $@. | agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | user-provided value |
4+
| agent_instructions.py:35:28:35:32 | ControlFlowNode for input | agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | agent_instructions.py:35:28:35:32 | ControlFlowNode for input | This prompt construction depends on a $@. | agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | user-provided value |
35
| openai_test.py:17:22:17:46 | ControlFlowNode for BinaryExpr | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | openai_test.py:17:22:17:46 | ControlFlowNode for BinaryExpr | This prompt construction depends on a $@. | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | user-provided value |
46
| openai_test.py:18:15:18:19 | ControlFlowNode for query | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | openai_test.py:18:15:18:19 | ControlFlowNode for query | This prompt construction depends on a $@. | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | user-provided value |
57
| openai_test.py:22:22:22:46 | ControlFlowNode for BinaryExpr | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | openai_test.py:22:22:22:46 | ControlFlowNode for BinaryExpr | This prompt construction depends on a $@. | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | user-provided value |
@@ -17,10 +19,16 @@
1719
edges
1820
| agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | agent_instructions.py:2:26:2:32 | ControlFlowNode for request | provenance | |
1921
| agent_instructions.py:2:26:2:32 | ControlFlowNode for request | agent_instructions.py:7:13:7:19 | ControlFlowNode for request | provenance | |
22+
| agent_instructions.py:2:26:2:32 | ControlFlowNode for request | agent_instructions.py:17:13:17:19 | ControlFlowNode for request | provenance | |
2023
| agent_instructions.py:7:5:7:9 | ControlFlowNode for input | agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | provenance | Sink:MaD:93 |
2124
| agent_instructions.py:7:13:7:19 | ControlFlowNode for request | agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
2225
| agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | provenance | dict.get |
2326
| agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | agent_instructions.py:7:5:7:9 | ControlFlowNode for input | provenance | |
27+
| agent_instructions.py:17:5:17:9 | ControlFlowNode for input | agent_instructions.py:25:28:25:32 | ControlFlowNode for input | provenance | |
28+
| agent_instructions.py:17:5:17:9 | ControlFlowNode for input | agent_instructions.py:35:28:35:32 | ControlFlowNode for input | provenance | |
29+
| agent_instructions.py:17:13:17:19 | ControlFlowNode for request | agent_instructions.py:17:13:17:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
30+
| agent_instructions.py:17:13:17:24 | ControlFlowNode for Attribute | agent_instructions.py:17:13:17:37 | ControlFlowNode for Attribute() | provenance | dict.get |
31+
| agent_instructions.py:17:13:17:37 | ControlFlowNode for Attribute() | agent_instructions.py:17:5:17:9 | ControlFlowNode for input | provenance | |
2432
| openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | openai_test.py:2:26:2:32 | ControlFlowNode for request | provenance | |
2533
| openai_test.py:2:26:2:32 | ControlFlowNode for request | openai_test.py:12:15:12:21 | ControlFlowNode for request | provenance | |
2634
| openai_test.py:2:26:2:32 | ControlFlowNode for request | openai_test.py:13:13:13:19 | ControlFlowNode for request | provenance | |
@@ -53,6 +61,12 @@ nodes
5361
| agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | semmle.label | ControlFlowNode for Attribute |
5462
| agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | semmle.label | ControlFlowNode for Attribute() |
5563
| agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | semmle.label | ControlFlowNode for BinaryExpr |
64+
| agent_instructions.py:17:5:17:9 | ControlFlowNode for input | semmle.label | ControlFlowNode for input |
65+
| agent_instructions.py:17:13:17:19 | ControlFlowNode for request | semmle.label | ControlFlowNode for request |
66+
| agent_instructions.py:17:13:17:24 | ControlFlowNode for Attribute | semmle.label | ControlFlowNode for Attribute |
67+
| agent_instructions.py:17:13:17:37 | ControlFlowNode for Attribute() | semmle.label | ControlFlowNode for Attribute() |
68+
| agent_instructions.py:25:28:25:32 | ControlFlowNode for input | semmle.label | ControlFlowNode for input |
69+
| agent_instructions.py:35:28:35:32 | ControlFlowNode for input | semmle.label | ControlFlowNode for input |
5670
| openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | semmle.label | ControlFlowNode for ImportMember |
5771
| openai_test.py:2:26:2:32 | ControlFlowNode for request | semmle.label | ControlFlowNode for request |
5872
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | semmle.label | ControlFlowNode for persona |

‎python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.qlref‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
1-
query: Security/CWE-1427/PromptInjection.ql
1+
query: experimental/Security/CWE-1427/PromptInjection.ql
22
postprocess: utils/test/InlineExpectationsTestQuery.ql

‎python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/agent_instructions.py‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,36 @@
33
app = Flask(__name__)
44

55
@app.route("/parameter-route")
6-
def get_input():
6+
def get_input1():
77
input = request.args.get("input")
88

99
agent = Agent(name="Assistant", instructions="This prompt is customized for " + input) # $Alert[py/prompt-injection]
1010

1111
result = Runner.run_sync(agent, "This is a user message.")
1212
print(result.final_output)
13+
14+
15+
@app.route("/parameter-route")
16+
def get_input2():
17+
input = request.args.get("input")
18+
19+
agent = Agent(name="Assistant", instructions="This prompt is not customized.")
20+
result = Runner.run_sync(
21+
agent=agent,
22+
input=[
23+
{
24+
"role": "user",
25+
"content": input, # $Alert[py/prompt-injection]
26+
}
27+
]
28+
)
29+
30+
result2 = Runner.run_sync(
31+
agent,
32+
[
33+
{
34+
"role": "user",
35+
"content": input, # $Alert[py/prompt-injection]
36+
}
37+
]
38+
)

0 commit comments

Comments
 (0)