52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
import asyncio
|
|
import sqlite3
|
|
from mcp.client.sse import sse_client
|
|
from mcp import ClientSession
|
|
|
|
def get_key(agent):
|
|
con = sqlite3.connect("data/context_hub.db")
|
|
cur = con.cursor()
|
|
cur.execute("SELECT api_key FROM api_keys WHERE agent_name=?", (agent,))
|
|
row = cur.fetchone()
|
|
return row[0]
|
|
|
|
async def call_as(agent, tool, args):
|
|
key = get_key(agent)
|
|
url = "http://localhost:8000/mcp"
|
|
headers = {"X-API-Key": key}
|
|
async with sse_client(url, headers=headers) as (read, write):
|
|
async with ClientSession(read, write) as session:
|
|
await session.initialize()
|
|
result = await session.call_tool(tool, args)
|
|
return result
|
|
|
|
async def main():
|
|
print("=== TEST 1 : GEMINI demande get_rules(scope=tt) -- doit etre refuse ===")
|
|
r = await call_as("GEMINI", "get_rules", {"scope": "tt"})
|
|
print(r.content[0].text)
|
|
print()
|
|
|
|
print("=== TEST 2 : GEMINI tente update_rule en se declarant CLAUDE dans les arguments (ancien vecteur) ===")
|
|
r = await call_as("GEMINI", "update_rule", {"scope": "tt", "agent": "CLAUDE", "data": {"test_spoofing": "FUITE SI VISIBLE"}})
|
|
print(r.content[0].text)
|
|
print()
|
|
|
|
print("=== TEST 3 (controle positif) : GEMINI accede a son propre scope llm -- doit reussir ===")
|
|
r = await call_as("GEMINI", "get_rules", {"scope": "llm"})
|
|
print(r.content[0].text[:150])
|
|
print()
|
|
|
|
print("=== TEST 4 (controle positif) : CLAUDE accede reellement a tt -- doit reussir ===")
|
|
r = await call_as("CLAUDE", "get_rules", {"scope": "tt"})
|
|
print(r.content[0].text[:150])
|
|
print()
|
|
|
|
print("=== TEST 5 : verification directe DB -- aucune ecriture spoofee dans tt ===")
|
|
con = sqlite3.connect("data/context_hub.db")
|
|
cur = con.cursor()
|
|
cur.execute("SELECT content FROM rules WHERE scope='tt'")
|
|
content = cur.fetchone()[0]
|
|
print("test_spoofing present dans tt ?", "test_spoofing" in content)
|
|
|
|
asyncio.run(main())
|