Skip to content

Commit e7a0fc7

Browse files
committed
python: Add query for prompt injection
This pull request introduces a new CodeQL query for detecting prompt injection vulnerabilities in Python code targeting AI prompting APIs such as agents and openai. The changes includes a new experimental query, new taint flow and type models, a customizable dataflow configuration, documentation, and comprehensive test coverage.
1 parent 34800d1 commit e7a0fc7

File tree

17 files changed

+519
-1
lines changed

17 files changed

+519
-1
lines changed

python/ql/integration-tests/query-suite/not_included_in_qls.expected

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ ql/python/ql/src/experimental/Security/CWE-079/EmailXss.ql
8787
ql/python/ql/src/experimental/Security/CWE-091/XsltInjection.ql
8888
ql/python/ql/src/experimental/Security/CWE-094/Js2Py.ql
8989
ql/python/ql/src/experimental/Security/CWE-1236/CsvInjection.ql
90+
ql/python/ql/src/experimental/Security/CWE-1427/PromptInjection.ql
9091
ql/python/ql/src/experimental/Security/CWE-176/UnicodeBypassValidation.ql
9192
ql/python/ql/src/experimental/Security/CWE-208/TimingAttackAgainstHash/PossibleTimingAttackAgainstHash.ql
9293
ql/python/ql/src/experimental/Security/CWE-208/TimingAttackAgainstHash/TimingAttackAgainstHash.ql
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
category: minorAnalysis
3+
---
4+
* Added experimental query `py/prompt-injection` to detect potential prompt injection vulnerabilities in code using LLMs.
5+
* Added taint flow model and type model for `agents` and `openai` modules.
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
extensions:
2+
- addsTo:
3+
pack: codeql/python-all
4+
extensible: sinkModel
5+
data:
6+
- ['agents', 'Member[Agent].Argument[instructions:]', 'prompt-injection']
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
extensions:
2+
- addsTo:
3+
pack: codeql/python-all
4+
extensible: sinkModel
5+
data:
6+
- ['OpenAI', 'Member[beta].Member[assistants].Member[create].Argument[instructions:]', 'prompt-injection']
7+
8+
- addsTo:
9+
pack: codeql/python-all
10+
extensible: typeModel
11+
data:
12+
- ['OpenAI', 'openai', 'Member[OpenAI,AsyncOpenAI,AzureOpenAI].ReturnValue']
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
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.</p>
9+
</overview>
10+
11+
<recommendation>
12+
<p>Sanitize user input and also avoid using user input in developer or system level prompts.</p>
13+
</recommendation>
14+
15+
<example>
16+
<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>
17+
<sample src="examples/example.py" />
18+
</example>
19+
20+
<references>
21+
<li>OpenAI: <a href="https://openai.github.io/openai-guardrails-python">Guardrails</a>.</li>
22+
</references>
23+
24+
</qhelp>
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
/**
2+
* @name Prompt injection
3+
* @kind path-problem
4+
* @problem.severity error
5+
* @security-severity 5.0
6+
* @precision high
7+
* @id py/prompt-injection
8+
* @tags security
9+
* experimental
10+
* external/cwe/cwe-1427
11+
*/
12+
13+
import python
14+
import experimental.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"
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
from flask import Flask, request
2+
from agents import Agent
3+
from guardrails import GuardrailAgent
4+
5+
@app.route("/parameter-route")
6+
def get_input():
7+
input = request.args.get("input")
8+
9+
goodAgent = GuardrailAgent( # GOOD: Agent created with guardrails automatically configured.
10+
config=Path("guardrails_config.json"),
11+
name="Assistant",
12+
instructions="This prompt is customized for " + input)
13+
14+
badAgent = Agent(
15+
name="Assistant",
16+
instructions="This prompt is customized for " + input # BAD: user input in agent instruction.
17+
)

python/ql/src/experimental/semmle/python/Concepts.qll

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -483,3 +483,28 @@ class EmailSender extends DataFlow::Node instanceof EmailSender::Range {
483483
*/
484484
DataFlow::Node getABody() { result in [super.getPlainTextBody(), super.getHtmlBody()] }
485485
}
486+
487+
/**
488+
* A data-flow node that prompts an AI model.
489+
*
490+
* Extend this class to refine existing API models. If you want to model new APIs,
491+
* extend `AIPrompt::Range` instead.
492+
*/
493+
class AIPrompt extends DataFlow::Node instanceof AIPrompt::Range {
494+
/** Gets an input that is used as AI prompt. */
495+
DataFlow::Node getAPrompt() { result = super.getAPrompt() }
496+
}
497+
498+
/** Provides a class for modeling new AI prompting mechanisms. */
499+
module AIPrompt {
500+
/**
501+
* A data-flow node that prompts an AI model.
502+
*
503+
* Extend this class to model new APIs. If you want to refine existing API models,
504+
* extend `AIPrompt` instead.
505+
*/
506+
abstract class Range extends DataFlow::Node {
507+
/** Gets an input that is used as AI prompt. */
508+
abstract DataFlow::Node getAPrompt();
509+
}
510+
}

python/ql/src/experimental/semmle/python/Frameworks.qll

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ private import experimental.semmle.python.frameworks.Scrapli
1313
private import experimental.semmle.python.frameworks.Twisted
1414
private import experimental.semmle.python.frameworks.JWT
1515
private import experimental.semmle.python.frameworks.Csv
16+
private import experimental.semmle.python.frameworks.OpenAI
1617
private import experimental.semmle.python.libraries.PyJWT
1718
private import experimental.semmle.python.libraries.Python_JWT
1819
private import experimental.semmle.python.libraries.Authlib
Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
/**
2+
* Provides classes modeling security-relevant aspects of the `openAI` Agents SDK package.
3+
* 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.
6+
*/
7+
8+
private import python
9+
private import semmle.python.ApiGraphs
10+
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.Runner` class. */
18+
API::Node classRef() { result = API::moduleImport("agents").getMember("Runner") }
19+
20+
/** Gets a reference to the `run` members. */
21+
API::Node runMembers() { result = classRef().getMember(["run", "run_sync", "run_streamed"]) }
22+
23+
/** Gets a reference to a potential property of `agents.Runner` called input which can refer to a system prompt depending on the role specified. */
24+
API::Node getContentNode() {
25+
result = runMembers().getKeywordParameter("input").getASubscript().getSubscript("content")
26+
or
27+
result = runMembers().getParameter(_).getASubscript().getSubscript("content")
28+
}
29+
}
30+
31+
/**
32+
* Provides models for Agent (instances of the `openai.OpenAI` class).
33+
*
34+
* See https://github.com/openai/openai-python.
35+
*/
36+
module OpenAI {
37+
/** Gets a reference to the `openai.OpenAI` class. */
38+
API::Node classRef() {
39+
result =
40+
API::moduleImport("openai").getMember(["OpenAI", "AsyncOpenAI", "AzureOpenAI"]).getReturn()
41+
}
42+
43+
/** Gets a reference to a potential property of `openai.OpenAI` called instructions which refers to the system prompt. */
44+
API::Node getContentNode() {
45+
exists(API::Node content |
46+
content =
47+
classRef()
48+
.getMember("responses")
49+
.getMember("create")
50+
.getKeywordParameter(["input", "instructions"])
51+
or
52+
content =
53+
classRef()
54+
.getMember("responses")
55+
.getMember("create")
56+
.getKeywordParameter(["input", "instructions"])
57+
.getASubscript()
58+
.getSubscript("content")
59+
or
60+
content =
61+
classRef()
62+
.getMember("realtime")
63+
.getMember("connect")
64+
.getReturn()
65+
.getMember("conversation")
66+
.getMember("item")
67+
.getMember("create")
68+
.getKeywordParameter("item")
69+
.getSubscript("content")
70+
or
71+
content =
72+
classRef()
73+
.getMember("chat")
74+
.getMember("completions")
75+
.getMember("create")
76+
.getKeywordParameter("messages")
77+
.getASubscript()
78+
.getSubscript("content")
79+
|
80+
// content
81+
if not exists(content.getASubscript())
82+
then result = content
83+
else
84+
// content.text
85+
result = content.getASubscript().getSubscript("text")
86+
)
87+
}
88+
}

0 commit comments

Comments
 (0)