diff --git a/python/ql/lib/semmle/python/frameworks/OpenAI.qll b/python/ql/lib/semmle/python/frameworks/OpenAI.qll index a247645081a4..ccb8f1bb6b6c 100644 --- a/python/ql/lib/semmle/python/frameworks/OpenAI.qll +++ b/python/ql/lib/semmle/python/frameworks/OpenAI.qll @@ -1,11 +1,32 @@ /** - * Provides classes modeling security-relevant aspects of the `openAI`Agents SDK package. + * Provides classes modeling security-relevant aspects of the `openAI` Agents SDK package. * See https://github.com/openai/openai-agents-python. + * As well as the regular openai python interface. + * See https://github.com/openai/openai-python. */ private import python private import semmle.python.ApiGraphs +/** + * Provides models for agents SDK (instances of the `agents.Runner` class etc). + * + * See https://github.com/openai/openai-agents-python. + */ +module AgentSDK { + /** Gets a reference to the `agents.Agent` class. */ + API::Node classRef() { result = API::moduleImport("agents").getMember("Runner") } + + API::Node runMembers() { result = classRef().getMember(["run", "run_sync", "run_streamed"]) } + + /** Gets a reference to a potential property of `agents.Runner` called input which can refer to a system prompt depending on the role specified. */ + API::Node getContentNode() { + result = runMembers().getKeywordParameter("input").getASubscript().getSubscript("content") + or + result = runMembers().getParameter(_).getASubscript().getSubscript("content") + } +} + /** * Provides models for Agent (instances of the `openai.OpenAI` class). * diff --git a/python/ql/lib/semmle/python/security/dataflow/PromptInjectionCustomizations.qll b/python/ql/lib/semmle/python/security/dataflow/PromptInjectionCustomizations.qll index c7f480d47b2e..aaa8c05418ed 100644 --- a/python/ql/lib/semmle/python/security/dataflow/PromptInjectionCustomizations.qll +++ b/python/ql/lib/semmle/python/security/dataflow/PromptInjectionCustomizations.qll @@ -52,6 +52,8 @@ module PromptInjection { private class PromptContentSink extends Sink { PromptContentSink() { this = OpenAI::getContentNode().asSink() + or + this = AgentSDK::getContentNode().asSink() } } diff --git a/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.expected b/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.expected index f52061f3c986..bbf9e5665bab 100644 --- a/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.expected +++ b/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.expected @@ -1,5 +1,7 @@ #select | 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 | +| 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 | +| 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 | | 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 | | 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 | | 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 @@ edges | agent_instructions.py:2:26:2:32 | ControlFlowNode for ImportMember | agent_instructions.py:2:26:2:32 | ControlFlowNode for request | provenance | | | agent_instructions.py:2:26:2:32 | ControlFlowNode for request | agent_instructions.py:7:13:7:19 | ControlFlowNode for request | provenance | | +| agent_instructions.py:2:26:2:32 | ControlFlowNode for request | agent_instructions.py:17:13:17:19 | ControlFlowNode for request | provenance | | | agent_instructions.py:7:5:7:9 | ControlFlowNode for input | agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | provenance | Sink:MaD:93 | | agent_instructions.py:7:13:7:19 | ControlFlowNode for request | agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep | | agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | provenance | dict.get | | agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | agent_instructions.py:7:5:7:9 | ControlFlowNode for input | provenance | | +| agent_instructions.py:17:5:17:9 | ControlFlowNode for input | agent_instructions.py:25:28:25:32 | ControlFlowNode for input | provenance | | +| agent_instructions.py:17:5:17:9 | ControlFlowNode for input | agent_instructions.py:35:28:35:32 | ControlFlowNode for input | provenance | | +| agent_instructions.py:17:13:17:19 | ControlFlowNode for request | agent_instructions.py:17:13:17:24 | ControlFlowNode for Attribute | provenance | AdditionalTaintStep | +| agent_instructions.py:17:13:17:24 | ControlFlowNode for Attribute | agent_instructions.py:17:13:17:37 | ControlFlowNode for Attribute() | provenance | dict.get | +| agent_instructions.py:17:13:17:37 | ControlFlowNode for Attribute() | agent_instructions.py:17:5:17:9 | ControlFlowNode for input | provenance | | | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | openai_test.py:2:26:2:32 | ControlFlowNode for request | provenance | | | openai_test.py:2:26:2:32 | ControlFlowNode for request | openai_test.py:12:15:12:21 | ControlFlowNode for request | provenance | | | 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 | agent_instructions.py:7:13:7:24 | ControlFlowNode for Attribute | semmle.label | ControlFlowNode for Attribute | | agent_instructions.py:7:13:7:37 | ControlFlowNode for Attribute() | semmle.label | ControlFlowNode for Attribute() | | agent_instructions.py:9:50:9:89 | ControlFlowNode for BinaryExpr | semmle.label | ControlFlowNode for BinaryExpr | +| agent_instructions.py:17:5:17:9 | ControlFlowNode for input | semmle.label | ControlFlowNode for input | +| agent_instructions.py:17:13:17:19 | ControlFlowNode for request | semmle.label | ControlFlowNode for request | +| agent_instructions.py:17:13:17:24 | ControlFlowNode for Attribute | semmle.label | ControlFlowNode for Attribute | +| agent_instructions.py:17:13:17:37 | ControlFlowNode for Attribute() | semmle.label | ControlFlowNode for Attribute() | +| agent_instructions.py:25:28:25:32 | ControlFlowNode for input | semmle.label | ControlFlowNode for input | +| agent_instructions.py:35:28:35:32 | ControlFlowNode for input | semmle.label | ControlFlowNode for input | | openai_test.py:2:26:2:32 | ControlFlowNode for ImportMember | semmle.label | ControlFlowNode for ImportMember | | openai_test.py:2:26:2:32 | ControlFlowNode for request | semmle.label | ControlFlowNode for request | | openai_test.py:12:5:12:11 | ControlFlowNode for persona | semmle.label | ControlFlowNode for persona | diff --git a/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.qlref b/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.qlref index 04d00b1b4b8c..08466562ffe7 100644 --- a/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.qlref +++ b/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/PromptInjection.qlref @@ -1,2 +1,2 @@ -query: Security/CWE-1427/PromptInjection.ql +query: experimental/Security/CWE-1427/PromptInjection.ql postprocess: utils/test/InlineExpectationsTestQuery.ql \ No newline at end of file diff --git a/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/agent_instructions.py b/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/agent_instructions.py index ebe26942ce43..12cebc1b5831 100644 --- a/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/agent_instructions.py +++ b/python/ql/test/experimental/query-tests/Security/CWE-1427-PromptInjection/agent_instructions.py @@ -3,10 +3,36 @@ app = Flask(__name__) @app.route("/parameter-route") -def get_input(): +def get_input1(): input = request.args.get("input") agent = Agent(name="Assistant", instructions="This prompt is customized for " + input) # $Alert[py/prompt-injection] result = Runner.run_sync(agent, "This is a user message.") print(result.final_output) + + +@app.route("/parameter-route") +def get_input2(): + input = request.args.get("input") + + agent = Agent(name="Assistant", instructions="This prompt is not customized.") + result = Runner.run_sync( + agent=agent, + input=[ + { + "role": "user", + "content": input, # $Alert[py/prompt-injection] + } + ] + ) + + result2 = Runner.run_sync( + agent, + [ + { + "role": "user", + "content": input, # $Alert[py/prompt-injection] + } + ] + )