End-to-End Encrypted Chat Completions

import os
import time
import requests
from typing import Optional
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

BASE_URL = os.environ.get("BASE_URL", "https://api.redpill.ai").rstrip("/")
API_KEY = os.environ["API_KEY"]
MODEL_NAME = os.environ.get("MODEL_NAME", "phala/gpt-oss-20b")
HKDF_INFO = b"ecdsa_encryption"


def to_uncompressed(pub_hex: str) -> bytes:
    b = bytes.fromhex(pub_hex)
    if len(b) == 64:
        return b"\x04" + b
    if len(b) == 65 and b[0] == 0x04:
        return b
    raise ValueError(f"bad pubkey length: {len(b)}")


def derive_key(shared: bytes) -> bytes:
    return HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=None,
        info=HKDF_INFO,
    ).derive(shared)


def encrypt_for_model(plaintext: str, model_pub_hex: str, aad: Optional[str]):
    model_pub = ec.EllipticCurvePublicKey.from_encoded_point(
        ec.SECP256K1(), to_uncompressed(model_pub_hex)
    )
    eph_priv = ec.generate_private_key(ec.SECP256K1())
    eph_pub = eph_priv.public_key()

    shared = eph_priv.exchange(ec.ECDH(), model_pub)
    key = derive_key(shared)
    nonce = os.urandom(12)
    aad_bytes = aad.encode("utf-8") if aad else None
    ct = AESGCM(key).encrypt(nonce, plaintext.encode("utf-8"), aad_bytes)

    eph_pub_bytes = eph_pub.public_bytes(
        encoding=serialization.Encoding.X962,
        format=serialization.PublicFormat.UncompressedPoint,
    )
    return (eph_pub_bytes + nonce + ct).hex()


def decrypt_from_model(enc_hex: str, client_priv, aad: Optional[str]):
    blob = bytes.fromhex(enc_hex)
    eph_pub = ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256K1(), blob[:65])
    nonce = blob[65:77]
    ct = blob[77:]

    shared = client_priv.exchange(ec.ECDH(), eph_pub)
    key = derive_key(shared)
    aad_bytes = aad.encode("utf-8") if aad else None

    pt = AESGCM(key).decrypt(nonce, ct, aad_bytes)
    return pt.decode("utf-8")


def client_pub_hex64(client_priv):
    pub = client_priv.public_key().public_bytes(
        encoding=serialization.Encoding.X962,
        format=serialization.PublicFormat.UncompressedPoint,
    )
    return pub[1:].hex()


def fetch_model_pubkey() -> str:
    r = requests.get(
        f"{BASE_URL}/v1/attestation/report",
        params={"model": MODEL_NAME, "signing_algo": "ecdsa"},
        headers={"Authorization": f"Bearer {API_KEY}"},
        timeout=60,
    )
    r.raise_for_status()
    data = r.json()
    key = data.get("signing_public_key")
    if not key:
        raise RuntimeError("signing_public_key not found in attestation response")
    return key


def main():
    model_pub_hex = fetch_model_pubkey()
    print(f"[OK] model pubkey length={len(model_pub_hex)}")

    client_priv = ec.generate_private_key(ec.SECP256K1())
    client_pub = client_pub_hex64(client_priv)

    nonce = "n" + str(int(time.time())) + "abcd1234abcd"
    ts = str(int(time.time()))
    prompt = "please only reply: REDPILL_E2EE_V2_OK"

    req_aad = f"v2|req|algo=ecdsa|model={MODEL_NAME}|m=0|c=-|n={nonce}|ts={ts}"
    enc_prompt = encrypt_for_model(prompt, model_pub_hex, req_aad)

    payload = {
        "model": MODEL_NAME,
        "stream": False,
        "messages": [{"role": "user", "content": enc_prompt}],
    }

    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json",
        "X-Signing-Algo": "ecdsa",
        "X-Client-Pub-Key": client_pub,
        "X-Model-Pub-Key": model_pub_hex,
        "X-E2EE-Version": "2",
        "X-E2EE-Nonce": nonce,
        "X-E2EE-Timestamp": ts,
    }

    r = requests.post(f"{BASE_URL}/v1/chat/completions", headers=headers, json=payload, timeout=120)
    r.raise_for_status()

    if r.headers.get("X-E2EE-Applied", "").lower() != "true":
        raise RuntimeError("X-E2EE-Applied is not true")

    data = r.json()
    enc_answer = data["choices"][0]["message"]["content"]

    resp_aad = (
        f"v2|resp|algo=ecdsa|model={data.get('model','')}|"
        f"id={data.get('id','')}|choice=0|field=content|n={nonce}|ts={ts}"
    )
    plain = decrypt_from_model(enc_answer, client_priv, resp_aad)
    print("decrypted:", plain)


if __name__ == "__main__":
    main()