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())