Skip to content

Commit 2c83dc3

Browse files
mbaludambaluda
committed
Use models as data
1 parent 04193f4 commit 2c83dc3

File tree

7 files changed

+63
-79
lines changed

7 files changed

+63
-79
lines changed

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,31 @@ private class EncodingAdditionalTaintStep extends TaintTracking::AdditionalTaint
325325
}
326326
}
327327

328+
/**
329+
* A data-flow node that prompts an AI model.
330+
*
331+
* Extend this class to refine existing API models. If you want to model new APIs,
332+
* extend `AIPrompt::Range` instead.
333+
*/
334+
class AIPrompt extends DataFlow::Node instanceof AIPrompt::Range {
335+
/** Gets an input that is used as AI prompt. */
336+
DataFlow::Node getAPrompt() { result = super.getAPrompt() }
337+
}
338+
339+
/** Provides a class for modeling new AI prompting mechanisms. */
340+
module AIPrompt {
341+
/**
342+
* A data-flow node that prompts an AI model.
343+
*
344+
* Extend this class to model new APIs. If you want to refine existing API models,
345+
* extend `AIPrompt` instead.
346+
*/
347+
abstract class Range extends DataFlow::Node {
348+
/** Gets an input that is logged. */
349+
abstract DataFlow::Node getAPrompt();
350+
}
351+
}
352+
328353
/**
329354
* A data-flow node that logs data.
330355
*

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@ 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
5857
private import semmle.python.frameworks.Opml
5958
private import semmle.python.frameworks.Oracledb
6059
private import semmle.python.frameworks.Pandas

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

Lines changed: 0 additions & 59 deletions
This file was deleted.

‎python/ql/lib/semmle/python/frameworks/agent.model.yml‎

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']

‎python/ql/lib/semmle/python/frameworks/openai.model.yml‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
extensions:
2+
- addsTo:
3+
pack: codeql/python-all
4+
extensible: sinkModel
5+
data:
6+
- ['OpenAI', 'Member[responses].Member[create].Argument[input:,instructions:]', 'prompt-injection']
7+
- ['OpenAI', 'Member[realtime].Member[connect].ReturnValue.Member[conversation].Member[item].Member[create].Argument[item:]', 'prompt-injection']
8+
- ['OpenAI', 'Member[chat].Member[completions].Member[create].Argument[messages:]', 'prompt-injection']
9+
10+
- addsTo:
11+
pack: codeql/python-all
12+
extensible: typeModel
13+
data:
14+
- ['OpenAI', 'openai', 'Member[OpenAI,AsyncOpenAI,AzureOpenAI].ReturnValue']

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ private import semmle.python.dataflow.new.DataFlow
99
private import semmle.python.Concepts
1010
private import semmle.python.dataflow.new.RemoteFlowSources
1111
private import semmle.python.dataflow.new.BarrierGuards
12-
private import semmle.python.frameworks.OpenAI
1312
private import semmle.python.frameworks.data.ModelsAsData
1413

1514
/**
@@ -39,10 +38,10 @@ module PromptInjection {
3938
private class ActiveThreatModelSourceAsSource extends Source, ActiveThreatModelSource { }
4039

4140
/**
42-
* Agent prompt sinks, considered as a flow sink.
41+
* A prompt to an AI model, considered as a flow sink.
4342
*/
44-
class SystemPromptSink extends Sink {
45-
SystemPromptSink() { this = [Agent::sink(), OpenAI::sink()].asSink() }
43+
class AIPromptAsSink extends Sink {
44+
AIPromptAsSink() { this = any(AIPrompt p).getAPrompt() }
4645
}
4746

4847
private class SinkFromModel extends Sink {

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,37 +12,37 @@
1212
edges
1313
| agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | agent_instructions.py:2:26:2:32 | ControlFlowNode for request | provenance | |
1414
| agent_instructions.py:2:26:2:32 | ControlFlowNode for request | agent_instructions.py:7:13:7:19 | ControlFlowNode for request | provenance | |
15-
| agent_instructions.py:7:5:7:9 | ControlFlowNode for input | agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | provenance | |
15+
| agent_instructions.py:7:5:7:9 | ControlFlowNode for input | agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | provenance | Sink:MaD:93 |
1616
| agent_instructions.py:7:13:7:19 | ControlFlowNode for request | agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
1717
| agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | provenance | dict.get |
1818
| agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | agent_instructions.py:7:5:7:9 | ControlFlowNode for input | provenance | |
1919
| openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | openai_test.py:2:26:2:32 | ControlFlowNode for request | provenance | |
2020
| openai_test.py:2:26:2:32 | ControlFlowNode for request | openai_test.py:12:15:12:21 | ControlFlowNode for request | provenance | |
2121
| openai_test.py:2:26:2:32 | ControlFlowNode for request | openai_test.py:13:13:13:19 | ControlFlowNode for request | provenance | |
2222
| openai_test.py:2:26:2:32 | ControlFlowNode for request | openai_test.py:14:12:14:18 | ControlFlowNode for request | provenance | |
23-
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:17:22:17:46 | ControlFlowNode for BinaryExpr | provenance | |
24-
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:22:22:22:46 | ControlFlowNode for BinaryExpr | provenance | |
25-
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:23:15:37:9 | ControlFlowNode for List | provenance | |
26-
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:41:22:41:46 | ControlFlowNode for BinaryExpr | provenance | |
27-
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:59:18:70:9 | ControlFlowNode for List | provenance | |
28-
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:74:18:83:9 | ControlFlowNode for List | provenance | |
23+
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:17:22:17:46 | ControlFlowNode for BinaryExpr | provenance | Sink:MaD:58613 |
24+
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:22:22:22:46 | ControlFlowNode for BinaryExpr | provenance | Sink:MaD:58613 |
25+
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:23:15:37:9 | ControlFlowNode for List | provenance | Sink:MaD:58613 |
26+
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:41:22:41:46 | ControlFlowNode for BinaryExpr | provenance | Sink:MaD:58613 |
27+
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:59:18:70:9 | ControlFlowNode for List | provenance | Sink:MaD:58615 |
28+
| openai_test.py:12:5:12:11 | ControlFlowNode for persona | openai_test.py:74:18:83:9 | ControlFlowNode for List | provenance | Sink:MaD:58615 |
2929
| openai_test.py:12:15:12:21 | ControlFlowNode for request | openai_test.py:12:15:12:26 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
3030
| openai_test.py:12:15:12:21 | ControlFlowNode for request | openai_test.py:13:13:13:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
3131
| openai_test.py:12:15:12:21 | ControlFlowNode for request | openai_test.py:14:12:14:23 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
3232
| openai_test.py:12:15:12:26 | ControlFlowNode for Attribute | openai_test.py:12:15:12:41 | ControlFlowNode for Attribute() | provenance | dict.get |
3333
| openai_test.py:12:15:12:41 | ControlFlowNode for Attribute() | openai_test.py:12:5:12:11 | ControlFlowNode for persona | provenance | |
34-
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:18:15:18:19 | ControlFlowNode for query | provenance | |
35-
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:23:15:37:9 | ControlFlowNode for List | provenance | |
36-
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:42:15:42:19 | ControlFlowNode for query | provenance | |
37-
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:47:18:55:13 | ControlFlowNode for Dict | provenance | |
38-
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:59:18:70:9 | ControlFlowNode for List | provenance | |
39-
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:74:18:83:9 | ControlFlowNode for List | provenance | |
34+
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:18:15:18:19 | ControlFlowNode for query | provenance | Sink:MaD:58613 |
35+
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:23:15:37:9 | ControlFlowNode for List | provenance | Sink:MaD:58613 |
36+
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:42:15:42:19 | ControlFlowNode for query | provenance | Sink:MaD:58613 |
37+
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:47:18:55:13 | ControlFlowNode for Dict | provenance | Sink:MaD:58614 |
38+
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:59:18:70:9 | ControlFlowNode for List | provenance | Sink:MaD:58615 |
39+
| openai_test.py:13:5:13:9 | ControlFlowNode for query | openai_test.py:74:18:83:9 | ControlFlowNode for List | provenance | Sink:MaD:58615 |
4040
| openai_test.py:13:13:13:19 | ControlFlowNode for request | openai_test.py:13:13:13:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
4141
| openai_test.py:13:13:13:19 | ControlFlowNode for request | openai_test.py:14:12:14:23 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
4242
| openai_test.py:13:13:13:24 | ControlFlowNode for Attribute | openai_test.py:13:13:13:37 | ControlFlowNode for Attribute() | provenance | dict.get |
4343
| openai_test.py:13:13:13:37 | ControlFlowNode for Attribute() | openai_test.py:13:5:13:9 | ControlFlowNode for query | provenance | |
44-
| openai_test.py:14:5:14:8 | ControlFlowNode for role | openai_test.py:47:18:55:13 | ControlFlowNode for Dict | provenance | |
45-
| openai_test.py:14:5:14:8 | ControlFlowNode for role | openai_test.py:59:18:70:9 | ControlFlowNode for List | provenance | |
44+
| openai_test.py:14:5:14:8 | ControlFlowNode for role | openai_test.py:47:18:55:13 | ControlFlowNode for Dict | provenance | Sink:MaD:58614 |
45+
| openai_test.py:14:5:14:8 | ControlFlowNode for role | openai_test.py:59:18:70:9 | ControlFlowNode for List | provenance | Sink:MaD:58615 |
4646
| openai_test.py:14:12:14:18 | ControlFlowNode for request | openai_test.py:14:12:14:23 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep |
4747
| openai_test.py:14:12:14:23 | ControlFlowNode for Attribute | openai_test.py:14:12:14:35 | ControlFlowNode for Attribute() | provenance | dict.get |
4848
| openai_test.py:14:12:14:35 | ControlFlowNode for Attribute() | openai_test.py:14:5:14:8 | ControlFlowNode for role | provenance | |

0 commit comments

Comments
 (0)