"""Canary leak check for an AI endpoint: does any answer contain another request's secret word?

Every request carries its own random canary (CANARY-XXXXXXXXXX) after a long shared prefix, so the
provider's prompt cache is in play. The model is asked to repeat its canary. If any answer contains a
canary that belongs to a different request, that's cross-request data in your responses.

A clean run doesn't prove a provider can't leak. It's a smoke alarm: cheap enough to run on a schedule.

Azure OpenAI:  AZURE_ENDPOINT=https://<res>.cognitiveservices.azure.com AZURE_KEY=... DEPLOYMENT=<name> \
               python3 canary_check.py --n 200 --workers 20
OpenAI-compatible (vLLM, SGLang, others):  BASE_URL=http://host:8000/v1 API_KEY=... MODEL=<model> python3 canary_check.py
Standard library only.
"""
import argparse, concurrent.futures as cf, json, os, random, re, string, time, urllib.error, urllib.request

PAT = re.compile(r"CANARY-[A-Z0-9]{10}")


def target():
    if os.environ.get("AZURE_ENDPOINT"):
        url = f"{os.environ['AZURE_ENDPOINT'].rstrip('/')}/openai/deployments/{os.environ['DEPLOYMENT']}/chat/completions?api-version=2025-01-01-preview"
        return url, {"api-key": os.environ["AZURE_KEY"]}, {}
    url = f"{os.environ['BASE_URL'].rstrip('/')}/chat/completions"
    return url, {"Authorization": f"Bearer {os.environ.get('API_KEY', '')}"}, {"model": os.environ["MODEL"]}


def shared_prefix(words):
    rnd = random.Random(7)  # identical for every request, so it gets cached and shared
    vocab = "refund invoice customer ticket order delivery warehouse policy clause manager approve return damaged".split()
    return "Support manual. " + " ".join(rnd.choice(vocab) for _ in range(words))


def one(url, headers, extra, prefix, canary, timeout):
    body = json.dumps({**extra, "temperature": 0, "max_tokens": 40, "messages": [
        {"role": "system", "content": prefix},
        {"role": "user", "content": f"My reference code is {canary}. Reply with my reference code and nothing else."}]}).encode()
    req = urllib.request.Request(url, data=body, method="POST", headers={**headers, "Content-Type": "application/json"})
    err = "retries"
    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r:
                d = json.load(r)
            text = d["choices"][0]["message"]["content"] or ""
            return {"canary": canary, "ok": True, "text": text, "cached": (d.get("usage", {}).get("prompt_tokens_details") or {}).get("cached_tokens")}
        except urllib.error.HTTPError as e:
            if e.code == 429:
                time.sleep(float(e.headers.get("retry-after") or 5))
                continue
            return {"canary": canary, "ok": False, "error": f"HTTP {e.code}"}
        except Exception as e:  # network errors: retry, then give up on this request
            time.sleep(2)
            err = type(e).__name__
    return {"canary": canary, "ok": False, "error": err}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=200, help="requests")
    ap.add_argument("--workers", type=int, default=20, help="requests in flight at once")
    ap.add_argument("--prefix-words", type=int, default=3000, help="shared prefix size (~1.3 tokens per word)")
    ap.add_argument("--out", default="canary_results.json")
    a = ap.parse_args()
    url, headers, extra = target()
    prefix = shared_prefix(a.prefix_words)
    canaries = ["CANARY-" + "".join(random.choices(string.ascii_uppercase + string.digits, k=10)) for _ in range(a.n)]
    t0 = time.time()
    with cf.ThreadPoolExecutor(a.workers) as ex:
        rows = list(ex.map(lambda c: one(url, headers, extra, prefix, c, 120), canaries))
    mine = set(canaries)
    ok = [r for r in rows if r["ok"]]
    for r in ok:
        found = set(PAT.findall(r["text"]))
        r["own_echoed"] = r["canary"] in found
        r["foreign"] = sorted((found & mine) - {r["canary"]})
    leaks = [r for r in ok if r["foreign"]]
    summary = {"requests": a.n, "answered": len(ok), "echoed_own_canary": sum(r["own_echoed"] for r in ok),
               "answers_with_cache_hits": sum(1 for r in ok if r.get("cached")), "answers_with_foreign_canary": len(leaks),
               "seconds": round(time.time() - t0, 1)}
    json.dump({"summary": summary, "rows": rows}, open(a.out, "w"), indent=1)
    print(json.dumps(summary, indent=1))
    print("LEAK FOUND" if leaks else "No foreign canaries found. (A clean run is a smoke alarm, not a proof.)")


if __name__ == "__main__":
    main()