#!/usr/bin/env python3
"""
DVeProto wire 1.0 — standalone reference (pip install cryptography).
Handshake: dve_hello / dve_client_ack (JSON, v=1). HKDF labels DVeProto-v1/c2s|s2c.
This file implements ONLY wire 1.0 after handshake.
"""
from __future__ import annotations
import base64, json, os, struct
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple, Union
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

PROTO_NAME = "DVeProto"
PROTO_HANDSHAKE_VERSION = 1
WIRE = "1.0"
INFO_C2S = b"DVeProto-v1/c2s"
INFO_S2C = b"DVeProto-v1/s2c"
NONCE_LEN = 12
GCM_TAG_LEN = 16

def _b64e(raw: bytes) -> str:
    return base64.b64encode(raw).decode("ascii")
def _b64d(s: str) -> bytes:
    return base64.b64decode(s.encode("ascii"), validate=True)
def derive_aes_keys(shared: bytes):
    c2s = HKDF(algorithm=hashes.SHA256(), length=32, salt=b"", info=INFO_C2S).derive(shared)
    s2c = HKDF(algorithm=hashes.SHA256(), length=32, salt=b"", info=INFO_S2C).derive(shared)
    return c2s, s2c

class DVeClientSession:
    def __init__(self, c2s: bytes, s2c: bytes):
        self._c2s, self._s2c = AESGCM(c2s), AESGCM(s2c)
        self.wire_version = WIRE
    @classmethod
    def from_server_hello_text(cls, hello_text: str):
        data = json.loads(hello_text)
        if data.get("type") != "dve_hello" or data.get("proto") != PROTO_NAME:
            raise ValueError("dve_hello expected")
        server_pub = X25519PublicKey.from_public_bytes(_b64d(data["server_pk"]))
        priv = X25519PrivateKey.generate()
        c2s, s2c = derive_aes_keys(priv.exchange(server_pub))
        pub = priv.public_key().public_bytes(encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw)
        ack = json.dumps({"type":"dve_client_ack","proto":PROTO_NAME,"v":1,"client_pk":_b64e(pub),"select":"1.0"}, separators=(",",":"))
        return cls(c2s, s2c), ack
    def pack_outgoing(self, obj: Dict[str, Any]) -> str:
        nonce = os.urandom(NONCE_LEN)
        pt = json.dumps(obj, separators=(",",":"), ensure_ascii=False).encode()
        ct = self._c2s.encrypt(nonce, pt, None)
        return json.dumps({"type":"dve","proto":PROTO_NAME,"v":1,"n":_b64e(nonce),"c":_b64e(ct)}, separators=(",",":"))
    def unpack_incoming(self, text: str) -> Dict[str, Any]:
        data = json.loads(text)
        pt = self._s2c.decrypt(_b64d(data["n"]), _b64d(data["c"]), None)
        return json.loads(pt.decode())

if __name__ == "__main__":
    print("DVeProto reference wire", WIRE)