Files
codeql/python/ql/test/query-tests/Security/CWE-1427-SystemPromptInjection/gemini_test.py
2026-07-02 15:42:18 +02:00

60 lines
1.6 KiB
Python

from google import genai
from google.genai import types
from flask import Flask, request # $ Source
app = Flask(__name__)
client = genai.Client()
@app.route("/gemini")
def get_input_gemini():
persona = request.args.get("persona")
query = request.args.get("query")
response1 = client.models.generate_content(
model="gemini-2.0-flash",
contents=[
{
"role": "model",
"parts": [
{
"text": "I am " + persona # $ Alert[py/system-prompt-injection]
}
]
},
{
"role": "user",
"parts": [
{
"text": query
}
]
}
],
config=types.GenerateContentConfig(
system_instruction="Talk like " + persona, # $ Alert[py/system-prompt-injection]
),
)
print(response1)
cache = client.caches.create(
model="gemini-2.0-flash",
config=types.CreateCachedContentConfig(
system_instruction="Talk like " + persona, # $ Alert[py/system-prompt-injection]
),
)
print(cache)
@app.route("/gemini-live")
async def get_input_gemini_live():
persona = request.args.get("persona")
async with client.aio.live.connect(
model="gemini-2.0-flash",
config=types.LiveConnectConfig(
system_instruction="Talk like " + persona, # $ Alert[py/system-prompt-injection]
),
) as session:
print(session)