Skip to content

Commit 005db5b

Browse files
knewbury01knewbury01
committed
Add first version prompt injection query python openai agents sdk
1 parent 5db6b92 commit 005db5b

File tree

9 files changed

+154
-0
lines changed

9 files changed

+154
-0
lines changed

‎python/ql/lib/semmle/python/Frameworks.qll‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ private import semmle.python.frameworks.Multidict
5454
private import semmle.python.frameworks.Mysql
5555
private import semmle.python.frameworks.MySQLdb
5656
private import semmle.python.frameworks.Numpy
57+
private import semmle.python.frameworks.OpenAI
5758
private import semmle.python.frameworks.Opml
5859
private import semmle.python.frameworks.Oracledb
5960
private import semmle.python.frameworks.Pandas

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
/**
2+
* Provides classes modeling security-relevant aspects of the `openAI`Agents SDK package.
3+
* See https://github.com/openai/openai-agents-python.
4+
*/
5+
6+
private import python
7+
private import semmle.python.ApiGraphs
8+
9+
/**
10+
* Provides models for Agent (instances of the `agents.Agent` class).
11+
*
12+
* See https://github.com/openai/openai-agents-python.
13+
*/
14+
module Agent {
15+
/** Gets a reference to the `agents.Agent` class. */
16+
API::Node classRef() { result = API::moduleImport("agents").getMember("Agent") }
17+
18+
/** Gets a reference to a potential property of `agents.Agent` called instructions which refers to the system prompt. */
19+
API::Node sink() { result = classRef().getACall().getKeywordParameter("instructions") }
20+
}

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
import python
2+
private import semmle.python.dataflow.new.DataFlow
3+
private import semmle.python.Concepts
4+
private import semmle.python.dataflow.new.RemoteFlowSources
5+
private import semmle.python.dataflow.new.BarrierGuards
6+
private import semmle.python.frameworks.OpenAI
7+
8+
/**
9+
* Provides default sources, sinks and sanitizers for detecting
10+
* "prompt injection"
11+
* vulnerabilities, as well as extension points for adding your own.
12+
*/
13+
module PromptInjection {
14+
/**
15+
* A data flow source for "prompt injection" vulnerabilities.
16+
*/
17+
abstract class Source extends DataFlow::Node { }
18+
19+
/**
20+
* A data flow sink for "prompt injection" vulnerabilities.
21+
*/
22+
abstract class Sink extends DataFlow::Node { }
23+
24+
/**
25+
* A sanitizer for "prompt injection" vulnerabilities.
26+
*/
27+
abstract class Sanitizer extends DataFlow::Node { }
28+
29+
/**
30+
* An active threat-model source, considered as a flow source.
31+
*/
32+
private class ActiveThreatModelSourceAsSource extends Source, ActiveThreatModelSource { }
33+
34+
class SystemPromptSink extends Sink {
35+
SystemPromptSink() { this = Agent::sink().asSink() }
36+
}
37+
}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
private import python
2+
import semmle.python.dataflow.new.DataFlow
3+
import semmle.python.dataflow.new.TaintTracking
4+
import PromptInjectionCustomizations::PromptInjection
5+
6+
private module PromptInjectionConfig implements DataFlow::ConfigSig {
7+
predicate isSource(DataFlow::Node node) { node instanceof Source }
8+
9+
predicate isSink(DataFlow::Node node) {
10+
node instanceof Sink
11+
//any()
12+
}
13+
14+
predicate isBarrierIn(DataFlow::Node node) { node instanceof Sanitizer }
15+
16+
predicate observeDiffInformedIncrementalMode() { any() }
17+
}
18+
19+
/** Global taint-tracking for detecting "prompt injection" vulnerabilities. */
20+
module PromptInjectionFlow = TaintTracking::Global<PromptInjectionConfig>;

‎python/ql/src/Security/CWE-1427/PromptInjection.qhelp‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
<!DOCTYPE qhelp PUBLIC
2+
"-//Semmle//qhelp//EN"
3+
"qhelp.dtd">
4+
<qhelp>
5+
6+
<overview>
7+
<p>Prompts can be constructed to bypass the original purposes of an agent and lead to sensitive data leak or
8+
operations that were not intended.
9+
</p>
10+
</overview>
11+
12+
<recommendation>
13+
Sanitize user input and also avoid using user input in developer or system level prompts.
14+
</recommendation>
15+
16+
<example>
17+
<p>In the following examples, the cases marked GOOD show secure prompt construction; whereas in the case marked BAD they may be susceptible to prompt injection.</p>
18+
<sample src="examples/TODO.py" />
19+
</example>
20+
21+
<references>
22+
<li>OWASP: <a href="https://owasp.org/www-community/attacks/PromptInjection">PromptInjection</a>.</li>
23+
</references>
24+
25+
</qhelp>

‎python/ql/src/Security/CWE-1427/PromptInjection.ql‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
/**
2+
* @name User input used in developer message and or system prompt
3+
* @description User input used in developer message and or system prompt can allow for Prompt Injection attacks.
4+
* @kind path-problem
5+
* @problem.severity error
6+
* @security-severity 5.0
7+
* @precision high
8+
* @id py/prompt-injection
9+
* @tags security
10+
* external/cwe/cwe-1427
11+
*/
12+
13+
import python
14+
import semmle.python.security.dataflow.PromptInjectionQuery
15+
import PromptInjectionFlow::PathGraph
16+
17+
from PromptInjectionFlow::PathNode source, PromptInjectionFlow::PathNode sink
18+
where PromptInjectionFlow::flowPath(source, sink)
19+
select sink.getNode(), source, sink, "This prompt construction depends on a $@.", source.getNode(),
20+
"user-provided value"

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
edges
2+
| agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | agent_instructions.py:2:26:2:32 | ControlFlowNode for request | provenance | |
3+
| agent_instructions.py:2:26:2:32 | ControlFlowNode for request | agent_instructions.py:7:13:7:19 | ControlFlowNode for request | provenance | |
4+
| agent_instructions.py:7:5:7:9 | ControlFlowNode for input | agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | provenance | |
5+
| agent_instructions.py:7:13:7:19 | ControlFlowNode for request | agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
6+
| agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | provenance | dict.get |
7+
| agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | agent_instructions.py:7:5:7:9 | ControlFlowNode for input | provenance | |
8+
nodes
9+
| agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | semmle.label | ControlFlowNode for ImportMember |
10+
| agent_instructions.py:2:26:2:32 | ControlFlowNode for request | semmle.label | ControlFlowNode for request |
11+
| agent_instructions.py:7:5:7:9 | ControlFlowNode for input | semmle.label | ControlFlowNode for input |
12+
| agent_instructions.py:7:13:7:19 | ControlFlowNode for request | semmle.label | ControlFlowNode for request |
13+
| agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | semmle.label | ControlFlowNode for Attribute |
14+
| agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | semmle.label | ControlFlowNode for Attribute() |
15+
| agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | semmle.label | ControlFlowNode for BinaryExpr |
16+
subpaths
17+
#select
18+
| 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 |

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
query: Security/CWE-1427/PromptInjection.ql

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from agents import Agent, Runner
2+
from flask import Flask, request # $ Source=flask
3+
app = Flask(__name__)
4+
5+
@app.route("/parameter-route")
6+
def get_input():
7+
input = request.args.get("input")
8+
9+
agent = Agent(name="Assistant", instructions="This prompt is customized for " + input) # $Alert[py/prompt-injection]
10+
11+
result = Runner.run_sync(agent, "This is a user message.")
12+
print(result.final_output)

0 commit comments

Comments
 (0)