import os, json, hmac, base64, hashlib, time from dataclasses import dataclass from typing import Dict, List, Optional, Tuple from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey from cryptography.hazmat.primitives.kdf.hkdf import HKDF from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes # ------------------------- Utils ------------------------- def b64e(b: bytes) -> str: return base64.b64encode(b).decode() def b64d(s: str) -> bytes: return base64.b64decode(s.encode()) def hkdf(ikm: bytes, info: bytes, length: int = 32) -> bytes: return HKDF(algorithm=hashes.SHA256(), length=length, salt=None, info=info).derive(ikm) def aes_cbc_enc(key: bytes, iv: bytes, pt: bytes) -> bytes: # PKCS#7 padding pad = 16 - (len(pt) % 16) pt = pt + bytes([pad])*pad cipher = Cipher(algorithms.AES(key), modes.CBC(iv)) return cipher.encryptor().update(pt) + cipher.encryptor().finalize() def aes_cbc_dec(key: bytes, iv: bytes, ct: bytes) -> bytes: cipher = Cipher(algorithms.AES(key), modes.CBC(iv)) pt = cipher.decryptor().update(ct) + cipher.decryptor().finalize() pad = pt[-1] if pad < 1 or pad > 16 or pt[-pad:] != bytes([pad])*pad: raise ValueError("Bad padding") return pt[:-pad] def hmac_sha256(key: bytes, data: bytes) -> bytes: return hmac.new(key, data, hashlib.sha256).digest() def x25519_shared(sk: X25519PrivateKey, pk: X25519PublicKey) -> bytes: return sk.exchange(pk) def pub_bytes(pk: X25519PublicKey) -> bytes: return pk.public_bytes(encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw) def priv_from_bytes(b: bytes) -> X25519PrivateKey: return X25519PrivateKey.from_private_bytes(b) def pub_from_bytes(b: bytes) -> X25519PublicKey: return X25519PublicKey.from_public_bytes(b) # ------------------ Data structures ---------------------- @dataclass class PublicBundle: identity_pub: str signed_prekey_pub: str signed_prekey_sig: str # 演示保留字段，未做真实签名校验 onetime_prekeys: Dict[str, str] # id -> pub @dataclass class SessionState: # 简化：只做单向发送链（对称 ratchet 计数） root_key: bytes chain_key: bytes counter: int peer_identity_pub_b64: str opk_id: Optional[str] # 用了哪个一次性预键（初始化时） established_at: float # ------------------------- Server ------------------------ class RelayServer: """ 极简“服务器”：存放用户公开密钥束，提供查询；转发密文帧。 不保存/解密消息正文。 """ def __init__(self): self.directory: Dict[str, PublicBundle] = {} self.mailbox: Dict[str, List[dict]] = {} def register_bundle(self, user: str, bundle: PublicBundle): self.directory[user] = bundle def fetch_bundle(self, user: str) -> Optional[PublicBundle]: return self.directory.get(user) def mark_onetime_used(self, user: str, opk_id: str): b = self.directory.get(user) if not b: return if opk_id in b.onetime_prekeys: del b.onetime_prekeys[opk_id] def send_frame(self, to_user: str, frame: dict): self.mailbox.setdefault(to_user, []).append(frame) def pull_frames(self, user: str) -> List[dict]: return self.mailbox.pop(user, []) # ------------------------- Client ------------------------ class Client: def __init__(self, name: str, server: RelayServer): self.name = name self.server = server # 长期身份密钥 self.ik_priv = X25519PrivateKey.generate() self.ik_pub = self.ik_priv.public_key() # 已签名预共享密钥（演示：不做真实签名） self.spk_priv = X25519PrivateKey.generate() self.spk_pub = self.spk_priv.public_key() self.spk_sig = os.urandom(64) # 占位 # 一次性预共享密钥池 self.opk_priv_map: Dict[str, X25519PrivateKey] = {} self.opk_pub_map: Dict[str, X25519PublicKey] = {} for i in range(5): sk = X25519PrivateKey.generate() pk = sk.public_key() opk_id = f"opk-{i}" self.opk_priv_map[opk_id] = sk self.opk_pub_map[opk_id] = pk # 会话（按对端用户） self.sessions: Dict[str, SessionState] = {} def publish(self): bundle = PublicBundle( identity_pub=b64e(pub_bytes(self.ik_pub)), signed_prekey_pub=b64e(pub_bytes(self.spk_pub)), signed_prekey_sig=b64e(self.spk_sig), onetime_prekeys={oid: b64e(pub_bytes(pk)) for oid, pk in self.opk_pub_map.items()} ) self.server.register_bundle(self.name, bundle) # ---------- 建会话（发起方）简化 X3DH ---------- def establish_session_as_initiator(self, peer: str) -> SessionState: pb = self.server.fetch_bundle(peer) assert pb, f"{peer} has no bundle" IKb = pub_from_bytes(b64d(pb.identity_pub)) SPKb = pub_from_bytes(b64d(pb.signed_prekey_pub)) # 选一个对方的一次性预键 opk_items = list(pb.onetime_prekeys.items()) opk_id, OPKb_b64 = opk_items[0] if opk_items else (None, None) OPKb = pub_from_bytes(b64d(OPKb_b64)) if OPKb_b64 else None # 发起者生成临时密钥对 EKa_priv = X25519PrivateKey.generate() EKa_pub = EKa_priv.public_key() # X3DH 4次 ECDH（无OPK则忽略第4项） s1 = x25519_shared(self.ik_priv, SPKb) s2 = x25519_shared(EKa_priv, IKb) s3 = x25519_shared(EKa_priv, SPKb) parts = [s1, s2, s3] if OPKb: s4 = x25519_shared(EKa_priv, OPKb) parts.append(s4) master_secret = b"".join(parts) root_key = hkdf(master_secret, info=b"ROOT", length=32) chain_key = hkdf(root_key, info=b"CHAIN", length=32) st = SessionState( root_key=root_key, chain_key=chain_key, counter=0, peer_identity_pub_b64=pb.identity_pub, opk_id=opk_id, established_at=time.time() ) self.sessions[peer] = st # 发送“建会话 + 第一条消息”的头信息（包含 EKa_pub、opk_id） self._pending_ephemeral_pub = b64e(pub_bytes(EKa_pub)) self._pending_opk_id = opk_id return st # ---------- 建会话（接收方） ---------- def _establish_session_as_responder(self, peer: str, ek_pub_b64: str, opk_id: Optional[str]) -> SessionState: EKa = pub_from_bytes(b64d(ek_pub_b64)) # 对方临时公钥 # 自己的密钥 IKb_priv = self.ik_priv SPKb_priv = self.spk_priv OPKb_priv = self.opk_priv_map.get(opk_id) if opk_id else None # 4次 ECDH（无OPK则忽略第4项） s1 = x25519_shared(SPKb_priv, pub_from_bytes(b64d(self.server.directory[peer].identity_pub))) # = ECDH(Ia, SPKb) # 注意：发起方 s1 是 ECDH(Ia, SPKb)，接收方等价项应是 ECDH(SPKb, Ia) # 但我们没有 Ia 私钥，这里换一种对称表达：按消息头与本地密钥构造相同串联 # 为确保与发起方一致，我们直接重算： # 对于接收方：ECDH(Iinitiator, Srecipient) == ECDH(Srecipient, Iinitiator) # 需要 Iinitiator 公钥：来自会话第一帧中？为简化，我们用目录中对方 identity_pub。 Ia_pub = pub_from_bytes(b64d(self.server.directory[peer].identity_pub)) s1 = x25519_shared(self.spk_priv, Ia_pub) s2 = x25519_shared(self.ik_priv, EKa) s3 = x25519_shared(self.spk_priv, EKa) parts = [s1, s2, s3] if OPKb_priv: s4 = x25519_shared(OPKb_priv, EKa) parts.append(s4) master_secret = b"".join(parts) root_key = hkdf(master_secret, info=b"ROOT", length=32) chain_key = hkdf(root_key, info=b"CHAIN", length=32) st = SessionState( root_key=root_key, chain_key=chain_key, counter=0, peer_identity_pub_b64=self.server.directory[peer].identity_pub, opk_id=opk_id, established_at=time.time() ) self.sessions[peer] = st # 一次性预键被使用后，服务端目录也标记删除 if opk_id: self.server.mark_onetime_used(self.name, opk_id) return st # ---------- 每条消息的派生与加密 ---------- def _derive_message_key(self, st: SessionState) -> Tuple[bytes, bytes, bytes]: """ 从 chain_key 派生本条消息的 (aes_key, hmac_key, iv)，然后更新 chain_key，counter+1 """ info = b"MSG|" + st.counter.to_bytes(8, "big") msg_key = hkdf(st.chain_key, info=info, length=80) # 32 AES + 32 HMAC + 16 IV aes_key = msg_key[:32] mac_key = msg_key[32:64] iv = msg_key[64:80] # 下一条链 st.chain_key = hkdf(st.chain_key, info=b"STEP", length=32) st.counter += 1 return aes_key, mac_key, iv # ---------- 发送 ---------- def send(self, to_user: str, plaintext: bytes): # 若没有会话，先建立 if to_user not in self.sessions: st = self.establish_session_as_initiator(to_user) ek_pub_b64 = self._pending_ephemeral_pub opk_id = self._pending_opk_id else: st = self.sessions[to_user] ek_pub_b64 = None opk_id = None aes_key, mac_key, iv = self._derive_message_key(st) ct = aes_cbc_enc(aes_key, iv, plaintext) mac = hmac_sha256(mac_key, iv + ct)[:8] frame = { "from": self.name, "to": to_user, "hdr": { "init": ek_pub_b64 is not None, "ek_pub_b64": ek_pub_b64, # 仅首次带上 "opk_id": opk_id, "counter": st.counter - 1, # 本条的计数 }, "body": { "iv": b64e(iv), "ct": b64e(ct), "mac": b64e(mac), } } self.server.send_frame(to_user, frame) # ---------- 接收 ---------- def receive_all(self) -> List[Tuple[str, bytes]]: frames = self.server.pull_frames(self.name) outputs = [] for f in frames: sender = f["from"] hdr = f["hdr"] body = f["body"] if sender not in self.sessions: # 首帧：建立被动会话 assert hdr["init"], "Missing init header" st = self._establish_session_as_responder( peer=sender, ek_pub_b64=hdr["ek_pub_b64"], opk_id=hdr["opk_id"] ) else: st = self.sessions[sender] # 按对方 counter 对齐（演示：假设顺序到达） aes_key, mac_key, iv = self._derive_message_key(st) if st.counter - 1 != hdr["counter"]: # 简化：严格顺序，真实实现需支持跳号、乱序恢复 raise ValueError("Out-of-order message (demo limitation)") if b64e(iv) != body["iv"]: # 教学演示：我们强制使用派生的 iv iv = b64d(body["iv"]) # 宽松一点也可以直接信任对端 iv ct = b64d(body["ct"]) mac = b64d(body["mac"]) calc = hmac_sha256(mac_key, iv + ct)[:8] if not hmac.compare_digest(mac, calc): raise ValueError("MAC verification failed") pt = aes_cbc_dec(aes_key, iv, ct) outputs.append((sender, pt)) return outputs # ------------------------- Demo -------------------------- def main(): server = RelayServer() alice = Client("alice", server) bob = Client("bob", server) # 发布各自的公钥包 alice.publish() bob.publish() # Alice 先发一条（会自动建会话） alice.send("bob", b"Hello Bob, this is Alice.") # Bob 拉取并解密 for frm in bob.receive_all(): print("[bob] got:", frm) # Bob 回复 bob.send("alice", b"Hi Alice, Bob here. Message received.") for frm in alice.receive_all(): print("[alice] got:", frm) # 连发多条，观察每条都换密钥（server 无法解密） for i in range(1, 4): alice.send("bob", f"Msg#{i} from Alice".encode()) for frm in bob.receive_all(): print("[bob] got:", frm) # 验证“服务器看不到明文” print("\n[server] directory keys (truncated):") print(json.dumps({ u: { "identity_pub": v.identity_pub[:24] + "...", "spk_pub": v.signed_prekey_pub[:24] + "...", "opk_count": len(v.onetime_prekeys) } for u, v in server.directory.items() }, indent=2)) if __name__ == "__main__": main()