diff --git a/.gitignore b/.gitignore index 2bca24a..011e332 100644 --- a/.gitignore +++ b/.gitignore @@ -18,4 +18,7 @@ __pycache__/ BACKUP/ #dont push test files yet -src/tests/* \ No newline at end of file +src/tests/* + +#dont add keys +src/keys/* \ No newline at end of file diff --git a/keys/sign-1.key b/keys/sign-1.key new file mode 100644 index 0000000..4832b41 --- /dev/null +++ b/keys/sign-1.key @@ -0,0 +1 @@ +&6ڵZH*qUg k mr \ No newline at end of file diff --git a/keys/sign-10.key b/keys/sign-10.key new file mode 100644 index 0000000..1109b2d --- /dev/null +++ b/keys/sign-10.key @@ -0,0 +1 @@ +yڿ$Ö!NWz:fX.! \ No newline at end of file diff --git a/keys/sign-2.key b/keys/sign-2.key new file mode 100644 index 0000000..d2019fd --- /dev/null +++ b/keys/sign-2.key @@ -0,0 +1 @@ +f}0J+Hz5i>RT2 \ No newline at end of file diff --git a/keys/sign-3.key b/keys/sign-3.key new file mode 100644 index 0000000..223ed26 --- /dev/null +++ b/keys/sign-3.key @@ -0,0 +1 @@ +nKAb`h{ұd8A \ No newline at end of file diff --git a/keys/sign-4.key b/keys/sign-4.key new file mode 100644 index 0000000..63ef524 --- /dev/null +++ b/keys/sign-4.key @@ -0,0 +1 @@ +Y+v:u@}s$Ἄ \ No newline at end of file diff --git a/keys/sign-5.key b/keys/sign-5.key new file mode 100644 index 0000000..49a1b0a --- /dev/null +++ b/keys/sign-5.key @@ -0,0 +1 @@ +H ^| 44E.mB=v \ No newline at end of file diff --git a/keys/sign-6.key b/keys/sign-6.key new file mode 100644 index 0000000..2b87688 Binary files /dev/null and b/keys/sign-6.key differ diff --git a/keys/sign-7.key b/keys/sign-7.key new file mode 100644 index 0000000..a572a43 --- /dev/null +++ b/keys/sign-7.key @@ -0,0 +1 @@ +iI$e;Z}De`J< \ No newline at end of file diff --git a/keys/sign-8.key b/keys/sign-8.key new file mode 100644 index 0000000..65b16b4 --- /dev/null +++ b/keys/sign-8.key @@ -0,0 +1 @@ +c$`d5hxjY@f \ No newline at end of file diff --git a/keys/sign-9.key b/keys/sign-9.key new file mode 100644 index 0000000..f9a6428 --- /dev/null +++ b/keys/sign-9.key @@ -0,0 +1 @@ +.L7%d6C!nW~faB㛣 \ No newline at end of file diff --git a/keys/verification_keys.json b/keys/verification_keys.json new file mode 100644 index 0000000..5d920a5 --- /dev/null +++ b/keys/verification_keys.json @@ -0,0 +1,12 @@ +{ + "1": "e95698fd8eaf11697932446eeb35a2f338c16187feaeaf781bde0fcf7a58b832", + "2": "75ef100757dafc45f2fb1e9b8f551753e11093ff338e786e7928b900372c543f", + "3": "f3ecde14e41eeb4124f3401f5ac0a25935aa806e890a131974fc6d4f052cae9c", + "4": "715630b3e16ee15835bcbc35fa8400718151f2f69a45a3b26bafc260772970b0", + "5": "9d5291c0905d39f4efb7ec23f3414dfee0dced314e3f11d40152b3030a927bb4", + "6": "335cb084d5f07f842225473583e68813c3fd11584f6844aa4d797c2393a38b78", + "7": "24fd00de107aab9752eb4b5db45ce1b930e823dcfbbf28786367e7e453c5e9e9", + "8": "8d61aacae21e996186c5d00c84f81cc880b355e2c55243adc8ba403a1fdda8e3", + "9": "ac18200dd846a212dfa111d75e4bfbc80ef28e202347b027ad5ebaa12515968b", + "10": "ba7bd8b99360ef69e6f53a337cd7f7bf598fffc7df1dd0bc559890abc6a6fec5" +} \ No newline at end of file diff --git a/run_all.sh b/run_all.sh index 3a448bf..6e442d0 100755 --- a/run_all.sh +++ b/run_all.sh @@ -13,36 +13,51 @@ mkdir -p "$LOG_DIR" PYTHON_BIN="$SRC_DIR/.venv/bin/python" if [[ ! -x "$PYTHON_BIN" ]]; then - PYTHON_BIN="python3" + PYTHON_BIN="python" fi + # Start server echo "Starting server, logging to $LOG_DIR/server.log..." "$PYTHON_BIN" "$SRC_DIR/server.py" > "$LOG_DIR/server.log" 2>&1 & SERVER_PID=$! + # Wait sleep 1 + +# Generate signing/verification keys +echo "Generating signing/verification keys..." +"$PYTHON_BIN" "$SRC_DIR/ttp.py" --N 10 + + # Start clients CLIENT_PIDS="" for i in 1 2 3 4 5 6 7 8 9 10; do - echo "Starting client $i, logging to $LOG_DIR/client${i}.log..." - "$PYTHON_BIN" "$SRC_DIR/client.py" --id "$i" --vec "1,2,4" > "$LOG_DIR/client${i}.log" 2>&1 & - CLIENT_PIDS="$CLIENT_PIDS $!" - sleep 0.1 + echo "Starting client $i, logging to $LOG_DIR/client${i}.log..." + "$PYTHON_BIN" "$SRC_DIR/client.py" --id "$i" --vec "1,2,4" --signingkey "keys/sign-${i}.key" --verificationkeys "keys/verification_keys.json" > "$LOG_DIR/client${i}.log" 2>&1 & + CLIENT_PIDS="$CLIENT_PIDS $!" + sleep 0.1 done + # Wait for clients to finish for client_pid in $CLIENT_PIDS; do - wait "$client_pid" || true - echo "Client $client_pid finished." - sleep 0.1 + wait "$client_pid" || true + echo "Client $client_pid finished." + sleep 0.1 done + # Stop server kill "$SERVER_PID" || true echo "Server $SERVER_PID stopped." + # Done echo "Done." + + + + diff --git a/src/_client_helper.py b/src/_client_helper.py index ba70319..5aa0431 100644 --- a/src/_client_helper.py +++ b/src/_client_helper.py @@ -1,7 +1,10 @@ #!/usr/bin/env python3 import json from cryptography.hazmat.primitives import serialization, hashes -from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey +from cryptography.hazmat.primitives.asymmetric.x25519 import ( + X25519PrivateKey, + X25519PublicKey, +) from cryptography.hazmat.primitives.kdf.hkdf import HKDF, HKDFExpand from cryptography.hazmat.primitives.ciphers.aead import AESGCM from Crypto.Random import get_random_bytes @@ -9,116 +12,169 @@ import requests import time import re -from config import DEBUG, DERIVED_KEY_LENGTH, MAX_POLLS, POLL_INTERVALS, DEBUG_TESTING_DELAY, DEBUG_TESTING_DELAY_TIME, DEBUG_TESTING_DELAY_CLIENT_ID, DEBUG_TESTING_DELAY_ROUND -from config import FIELD_ELEMENT_SIZE, R +from config import ( + DEBUG, + DERIVED_KEY_LENGTH, + MAX_POLLS, + POLL_INTERVALS, + DEBUG_TESTING_DELAY, + DEBUG_TESTING_DELAY_TIME, + DEBUG_TESTING_DELAY_CLIENT_ID, + DEBUG_TESTING_DELAY_ROUND, +) +from config import FIELD_ELEMENT_SIZE, R + def bytes_to_field_element(b: bytes) -> int: - return int.from_bytes(b, byteorder='big') % R + return int.from_bytes(b, byteorder="big") % R + def field_elements_to_bytes(x: int) -> bytes: - return (x % R).to_bytes(FIELD_ELEMENT_SIZE, byteorder='big') + return (x % R).to_bytes(FIELD_ELEMENT_SIZE, byteorder="big") + + +def field_negate(x: int) -> int: + return (R - x) % R + -def field_negate(x : int) -> int: - return (R-x) % R +def field_add(a: int, b: int) -> int: + return (a + b) % R -def field_add(a:int, b:int) -> int: - return (a+b) % R def prg_block_to_field_elements(prg_block: bytes, vec_len: int) -> list[int]: - return [bytes_to_field_element(prg_block[j*FIELD_ELEMENT_SIZE: (j+1)*FIELD_ELEMENT_SIZE]) for j in range(vec_len)] + return [ + bytes_to_field_element( + prg_block[j * FIELD_ELEMENT_SIZE : (j + 1) * FIELD_ELEMENT_SIZE] + ) + for j in range(vec_len) + ] + + +def pubkey_to_bytes(pubkey: X25519PublicKey) -> bytes: + return pubkey.public_bytes( + encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw + ) + def pubkey_to_b64(pubkey: X25519PublicKey) -> str: pubkey_raw = pubkey.public_bytes( - encoding=serialization.Encoding.Raw, - format=serialization.PublicFormat.Raw + encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw ) - return base64.b64encode(pubkey_raw).decode('ascii') + return base64.b64encode(pubkey_raw).decode("ascii") + def b64_to_pubkey(b64_str: str) -> X25519PublicKey: - pubkey_decoded = base64.b64decode(b64_str.encode('ascii')) + pubkey_decoded = base64.b64decode(b64_str.encode("ascii")) return X25519PublicKey.from_public_bytes(pubkey_decoded) + def privkey_to_raw_bytes(privkey: X25519PrivateKey) -> bytes: privkey_raw = privkey.private_bytes( encoding=serialization.Encoding.Raw, format=serialization.PrivateFormat.Raw, - encryption_algorithm=serialization.NoEncryption() # do not encrypt the key for debugging + encryption_algorithm=serialization.NoEncryption(), # do not encrypt the key for debugging ) return privkey_raw + def privkey_to_b64(privkey: X25519PrivateKey) -> str: privkey_raw = privkey_to_raw_bytes(privkey) - return base64.b64encode(privkey_raw).decode('ascii') + return base64.b64encode(privkey_raw).decode("ascii") + def b64_to_privkey(b64_str: str) -> X25519PrivateKey: - privkey_decoded = base64.b64decode(b64_str.encode('ascii')) + privkey_decoded = base64.b64decode(b64_str.encode("ascii")) return X25519PrivateKey.from_private_bytes(privkey_decoded) + def _poll_for_round1_result(client_id: int, server_url: str) -> dict: - """Poll server for round 1 result. Result response from server, or None if timeout""" - print(f"Client {client_id}: Polling for round 1 result...", flush=True) - for poll_count in range(MAX_POLLS[1]): - try: - result_resp = requests.get(f"{server_url}/round1/result?client_id={client_id}", timeout=5) - if result_resp.status_code == 200: - print(f"Client {client_id}: Got round 1 result after {poll_count} polls", flush=True) - return result_resp.json() - except requests.exceptions.Timeout: - pass # Continue polling - except requests.exceptions.RequestException as e: - print(f"Client {client_id}: Error polling: {e}", flush=True) - - time.sleep(POLL_INTERVALS[1]) - - print(f"Client {client_id}: Timeout waiting for round 1 result", flush=True) - return None + """Poll server for round 1 result. Result response from server, or None if timeout""" + print(f"Client {client_id}: Polling for round 1 result...", flush=True) + for poll_count in range(MAX_POLLS[1]): + try: + result_resp = requests.get( + f"{server_url}/round1/result?client_id={client_id}", timeout=5 + ) + if result_resp.status_code == 200: + print( + f"Client {client_id}: Got round 1 result after {poll_count} polls", + flush=True, + ) + return result_resp.json() + except requests.exceptions.Timeout: + pass # Continue polling + except requests.exceptions.RequestException as e: + print(f"Client {client_id}: Error polling: {e}", flush=True) + + time.sleep(POLL_INTERVALS[1]) + + print(f"Client {client_id}: Timeout waiting for round 1 result", flush=True) + return None + def _poll_for_round2_result(client_id: int, server_url: str) -> dict: - """Poll server for round 2 result""" - print(f"Client {client_id}: Polling for round 2 result...", flush=True) - for poll_count in range(MAX_POLLS[2]): - try: - result_resp = requests.get(f"{server_url}/round2/result?client_id={client_id}", timeout=5) - if result_resp.status_code == 200: - print(f"Client {client_id}: Got round 2 result after {poll_count} polls", flush=True) - return result_resp.json() - except requests.exceptions.Timeout: - pass # Continue polling - except requests.exceptions.RequestException as e: - print(f"Client {client_id}: Error polling: {e}", flush=True) - - time.sleep(POLL_INTERVALS[2]) - - print(f"Client {client_id}: Timeout waiting for round 2 result", flush=True) - return None + """Poll server for round 2 result""" + print(f"Client {client_id}: Polling for round 2 result...", flush=True) + for poll_count in range(MAX_POLLS[2]): + try: + result_resp = requests.get( + f"{server_url}/round2/result?client_id={client_id}", timeout=5 + ) + if result_resp.status_code == 200: + print( + f"Client {client_id}: Got round 2 result after {poll_count} polls", + flush=True, + ) + return result_resp.json() + except requests.exceptions.Timeout: + pass # Continue polling + except requests.exceptions.RequestException as e: + print(f"Client {client_id}: Error polling: {e}", flush=True) + + time.sleep(POLL_INTERVALS[2]) + + print(f"Client {client_id}: Timeout waiting for round 2 result", flush=True) + return None + def _poll_for_round3_result(client_id: int, server_url: str) -> dict: """Poll server for round 3 result""" print(f"Client {client_id}: Polling for round 3 result...", flush=True) for poll_count in range(MAX_POLLS[3]): try: - result_resp = requests.get(f"{server_url}/round3/result?client_id={client_id}", timeout=5) + result_resp = requests.get( + f"{server_url}/round3/result?client_id={client_id}", timeout=5 + ) if result_resp.status_code == 200: - print(f"Client {client_id}: Got round 3 result after {poll_count} polls", flush=True) + print( + f"Client {client_id}: Got round 3 result after {poll_count} polls", + flush=True, + ) return result_resp.json() except requests.exceptions.Timeout: pass # Continue polling except requests.exceptions.RequestException as e: print(f"Client {client_id}: Error polling: {e}", flush=True) - + time.sleep(POLL_INTERVALS[3]) - + print(f"Client {client_id}: Timeout waiting for round 3 result", flush=True) return None + def _poll_for_round4_result(client_id: int, server_url: str) -> dict: """Poll server for round 4 result""" print(f"Client {client_id}: Polling for round 4 result...", flush=True) for poll_count in range(MAX_POLLS[4]): try: - result_resp = requests.get(f"{server_url}/round4/result?client_id={client_id}", timeout=5) + result_resp = requests.get( + f"{server_url}/round4/result?client_id={client_id}", timeout=5 + ) if result_resp.status_code == 200: - print(f"Client {client_id}: Got round 4 result after {poll_count} polls", flush=True) + print( + f"Client {client_id}: Got round 4 result after {poll_count} polls", + flush=True, + ) return result_resp.json() except requests.exceptions.Timeout: pass # Continue polling @@ -128,68 +184,134 @@ def _poll_for_round4_result(client_id: int, server_url: str) -> dict: print(f"Client {client_id}: Timeout waiting for round 4 result", flush=True) return None + +def _poll_for_round5_result(client_id: int, server_url: str) -> dict: + """Poll server for round 5 result""" + print(f"Client {client_id}: Polling for round 5 result...", flush=True) + for poll_count in range(MAX_POLLS[5]): + try: + result_resp = requests.get( + f"{server_url}/round5/result?client_id={client_id}", timeout=5 + ) + if result_resp.status_code == 200: + print( + f"Client {client_id}: Got round 5 result after {poll_count} polls", + flush=True, + ) + return result_resp.json() + except requests.exceptions.Timeout: + pass + except requests.exceptions.RequestException as e: + print(f"Client {client_id}: Error polling: {e}", flush=True) + time.sleep(POLL_INTERVALS[5]) + print(f"Client {client_id}: Timeout waiting for round 5 result", flush=True) + return None + + def do_round(client_id: int, round: int, payload: dict, server_url: str) -> dict: - if DEBUG_TESTING_DELAY and client_id == DEBUG_TESTING_DELAY_CLIENT_ID and round == DEBUG_TESTING_DELAY_ROUND: - print(f"Client {client_id}: Delaying before sending round {round}...", flush=True) + if ( + DEBUG_TESTING_DELAY + and client_id == DEBUG_TESTING_DELAY_CLIENT_ID + and round == DEBUG_TESTING_DELAY_ROUND + ): + print( + f"Client {client_id}: Delaying before sending round {round}...", flush=True + ) time.sleep(DEBUG_TESTING_DELAY_TIME) print(f"Client {client_id}: Sending round {round} payload: {payload}", flush=True) round_resp = requests.post(f"{server_url}/round/{round}", json=payload) - print(f"Client {client_id}: Round {round} immediate response: {round_resp.json()}", flush=True) + print( + f"Client {client_id}: Round {round} immediate response: {round_resp.json()}", + flush=True, + ) result = poll_for_round_result(client_id, round, server_url) return result + def poll_for_round_result(client_id: int, round: int, server_url: str) -> dict: if round == 1: - return _poll_for_round1_result(client_id, server_url) + return _poll_for_round1_result(client_id, server_url) elif round == 2: - return _poll_for_round2_result(client_id, server_url) + return _poll_for_round2_result(client_id, server_url) elif round == 3: - return _poll_for_round3_result(client_id, server_url) + return _poll_for_round3_result(client_id, server_url) elif round == 4: - return _poll_for_round4_result(client_id, server_url) + return _poll_for_round4_result(client_id, server_url) + elif round == 5: + return _poll_for_round5_result(client_id, server_url) else: - print(f"Client {client_id}: No polling implemented for round {round}", flush=True) - return {} - - -def derive_shared_key(client_id: int, other_id: int, self_c_sec: X25519PrivateKey, other_c_pub: X25519PublicKey) -> bytes: + print( + f"Client {client_id}: No polling implemented for round {round}", flush=True + ) + return {} + + +def derive_shared_key( + client_id: int, + other_id: int, + self_c_sec: X25519PrivateKey, + other_c_pub: X25519PublicKey, +) -> bytes: shared_key = self_c_sec.exchange(other_c_pub) id_low, id_high = min(client_id, other_id), max(client_id, other_id) derived_key = HKDF( algorithm=hashes.SHA256(), length=DERIVED_KEY_LENGTH, salt=None, - info=b'key-agreement-pair-'+str(id_low).encode('ascii')+b'-'+str(id_high).encode('ascii'), + info=b"key-agreement-pair-" + + str(id_low).encode("ascii") + + b"-" + + str(id_high).encode("ascii"), ).derive(shared_key) return derived_key -def make_prg(client_id: int, other_id: int, self_s_sec: X25519PrivateKey, other_s_pub: X25519PublicKey, vec_len: int) -> bytes: - shared_key = self_s_sec.exchange(other_s_pub) - id_low, id_high = min(client_id, other_id), max(client_id, other_id) - return HKDFExpand(algorithm=hashes.SHA256(), length=vec_len*FIELD_ELEMENT_SIZE, info=b'prg-pair-'+str(id_low).encode('ascii')+b'-'+str(id_high).encode - ('ascii')).derive(shared_key) - -def make_prg2(client_id: int, other_id:int, prg_seed: bytes, vec_len:int) -> bytes: - PRG_block = HKDF( - algorithm=hashes.SHA256(), - length=vec_len*FIELD_ELEMENT_SIZE, - salt=None, - info=b'prg-seed-self'+str(client_id).encode('ascii')+b'-other'+str(other_id).encode('ascii'), - ).derive(prg_seed) - return PRG_block - -def encrypt_with_derived_key(derived_key: bytes, plaintext: bytes, - associated_data: bytes) -> tuple[bytes, bytes]: + +def make_prg( + client_id: int, + other_id: int, + self_s_sec: X25519PrivateKey, + other_s_pub: X25519PublicKey, + vec_len: int, +) -> bytes: + shared_key = self_s_sec.exchange(other_s_pub) + id_low, id_high = min(client_id, other_id), max(client_id, other_id) + return HKDFExpand( + algorithm=hashes.SHA256(), + length=vec_len * FIELD_ELEMENT_SIZE, + info=b"prg-pair-" + + str(id_low).encode("ascii") + + b"-" + + str(id_high).encode("ascii"), + ).derive(shared_key) + + +def make_prg2(client_id: int, other_id: int, prg_seed: bytes, vec_len: int) -> bytes: + PRG_block = HKDF( + algorithm=hashes.SHA256(), + length=vec_len * FIELD_ELEMENT_SIZE, + salt=None, + info=b"prg-seed-self" + + str(client_id).encode("ascii") + + b"-other" + + str(other_id).encode("ascii"), + ).derive(prg_seed) + return PRG_block + + +def encrypt_with_derived_key( + derived_key: bytes, plaintext: bytes, associated_data: bytes +) -> tuple[bytes, bytes]: aesgcm = AESGCM(derived_key) - nonce = get_random_bytes(12) # 96-bit nonce for AES-GCM + nonce = get_random_bytes(12) # 96-bit nonce for AES-GCM ciphertext = aesgcm.encrypt(nonce, plaintext, associated_data) return nonce, ciphertext -def decrypt_with_derived_key(derived_key: bytes, nonce: bytes, ciphertext: bytes, - associated_data) -> bytes: +def decrypt_with_derived_key( + derived_key: bytes, nonce: bytes, ciphertext: bytes, associated_data +) -> bytes: """ Decrypt ciphertext with the same AES-GCM key and nonce. Raises InvalidTag if authentication fails. @@ -198,44 +320,62 @@ def decrypt_with_derived_key(derived_key: bytes, nonce: bytes, ciphertext: bytes plaintext = aesgcm.decrypt(nonce, ciphertext, associated_data) return plaintext + def ids_to_associated_data(from_id: int, to_id: int) -> bytes: - return f"from:{from_id},to:{to_id}".encode('ascii') + return f"from:{from_id},to:{to_id}".encode("ascii") + def associated_data_to_ids(associated_data: bytes) -> tuple[int, int]: - match = re.search(r"from:([0-9]+),to:([0-9]+)", associated_data.decode('ascii')) + match = re.search(r"from:([0-9]+),to:([0-9]+)", associated_data.decode("ascii")) if match: return int(match.group(1)), int(match.group(2)) - raise ValueError(f"Invalid associated_data format, unable to extract from and to from {associated_data.decode('ascii')}") + raise ValueError( + f"Invalid associated_data format, unable to extract from and to from {associated_data.decode('ascii')}" + ) + def bencode(data: bytes) -> str: - return base64.b64encode(data).decode('ascii') + return base64.b64encode(data).decode("ascii") + def bdecode_to_bytes(data_b64: str) -> bytes: - return base64.b64decode(data_b64.encode('ascii')) + return base64.b64decode(data_b64.encode("ascii")) + def bdecode(data_b64: str) -> bytes: return bdecode_to_bytes(data_b64) + def bdecode_to_str(data_b64: str) -> str: - return base64.b64decode(data_b64.encode('ascii')).decode('utf-8') + return base64.b64decode(data_b64.encode("ascii")).decode("utf-8") + def jencode_to_bytes(json_data: dict) -> bytes: - return json.dumps(json_data).encode('utf-8') + return json.dumps(json_data).encode("utf-8") + + +def bytes_to_json(data_bytes: bytes) -> dict: + return json.loads(data_bytes.decode("utf-8")) -def bytes_to_json(data_bytes: bytes) -> dict: - return json.loads(data_bytes.decode('utf-8')) def jdecode_from_bytes(data_bytes: bytes) -> dict: return bytes_to_json(data_bytes) + def jencode_to_b64str(json_data: dict) -> str: - return bencode(json.dumps(json_data).encode('utf-8')) + return bencode(json.dumps(json_data).encode("utf-8")) + -def b64str_to_json(data_b64_str: str) -> dict: +def b64str_to_json(data_b64_str: str) -> dict: return bytes_to_json(bdecode_to_bytes(data_b64_str)) + def jencode_ciphertexts_for_other_r1r_r2(ciphertexts_for_other_r1r_r2: dict) -> dict: - return {r1r: (bencode(nonce), bencode(ciphertext)) for (r1r, (nonce, ciphertext)) in ciphertexts_for_other_r1r_r2.items()} # Placeholder + return { + r1r: (bencode(nonce), bencode(ciphertext)) + for (r1r, (nonce, ciphertext)) in ciphertexts_for_other_r1r_r2.items() + } # Placeholder + -def jdecode_ciphertexts(nonce_b64, ciphertext_b64) -> tuple[bytes,bytes]: - return (bdecode(nonce_b64), bdecode(ciphertext_b64)) \ No newline at end of file +def jdecode_ciphertexts(nonce_b64, ciphertext_b64) -> tuple[bytes, bytes]: + return (bdecode(nonce_b64), bdecode(ciphertext_b64)) diff --git a/src/_server_helper.py b/src/_server_helper.py index 34a3cea..ed306e1 100644 --- a/src/_server_helper.py +++ b/src/_server_helper.py @@ -1,141 +1,179 @@ #!/usr/bin/env python3 from flask import jsonify -from _client_helper import b64_to_pubkey, bdecode, field_add, field_negate, make_prg, make_prg2, prg_block_to_field_elements +from _client_helper import ( + b64_to_pubkey, + bdecode, + field_add, + field_negate, + make_prg, + make_prg2, + prg_block_to_field_elements, +) from config import ROUNDS, DEBUG, MAX_CLIENTS, THRESHOLD_CLIENTS from Crypto.Protocol.SecretSharing import Shamir from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey -def extract_round_client_id_payload(data: dict, expected_round: int) -> tuple[int, int, dict]: - # Extract round number from JSON body - round = data.get('round') - # Handle errors - if round is None: - raise ValueError('Error: Missing "round" field in JSON body') - round = int(round) - assert round == expected_round, f"Round number {expected_round} from POST does not match JSON body round number {round}" - if not (1 <= round <= ROUNDS): - raise ValueError(f"Error: Invalid round {round}; must be 1..{ROUNDS}") - - - # Extract client_id from JSON body - client_id = data.get('client_id') - # TODO V2: HANDLE AUTHENTICATION! CURRENTLY WE DO NOT CHECK TO ENSURE CLIENT NOT LYING ABOUT ID. - # Handle errors - if client_id is None: - raise ValueError('Error: Missing "client_id" field in JSON body') - client_id = int(client_id) - if not (1 <= client_id <= MAX_CLIENTS): - raise ValueError(f"Error: Invalid client_id {client_id}; must be 1..{MAX_CLIENTS}") - - # Extract payload from JSON body - payload = data.get('payload') - - return (round, client_id, payload) - - - - -def build_keyset_response(received_data: dict[int, dict[int, dict]], round1_responders: set[int]) -> dict[int, dict]: - # received_data format: {round: {client_id: {payload}}} - round1_data = received_data[1] - return {client_id: round1_data[client_id] for client_id in round1_responders} - - -def build_sharekeys_response(client_id: int, received_data: dict[int, dict[int, dict]], round2_responders: set[int]) -> dict[int, dict]: - # received_data format: {round: {client_id: {payload}}} - # Received from client 3 in round 2: {'1': ['jNrWa3PZPHDvX391', 'xxx'], '2': ['kDUON98rzqdA98+n', 'xxx']} - # TODO: CHANGE RECEIVED DATA TO ONLY INPUTTING THIS ROUND'S DATA - round2_data = received_data[2] - to_ret = {} - for other_id in round2_responders: - if other_id == client_id: - continue - if other_id not in round2_data or str(client_id) not in round2_data[other_id]: - raise ValueError(f"Error: build_sharekeys_response called for client {client_id}, but no round 2 data from other client {other_id} for this client") - to_ret[other_id] = round2_data[other_id][str(client_id)] # add received nonce and ciphertext from other_id for client_id to client_id's dictionary - return to_ret - - -def build_masked_input_response(client_id: int, received_data: dict[int, dict[int, dict]], round3_responders: set[int]) -> dict[int, dict]: - return {str(uid): {} for uid in round3_responders} - - -def compute_final_aggregate(received_data, round2_responders, round3_responders, round4_data, vec_len): - """ - Compute the final aggregate: z = sum(y_u) - sum(p_u) - sum(p_{s,d}) - where s are surviving users and d are dropped users. - """ - survived = set(round3_responders) - dropped = set(round2_responders) - survived - - - # Step 1: Sum all y_u for u in U3 - z = [0] * vec_len - for u in round3_responders: - y_u = received_data[3][u] - for j in range(vec_len): - z[j] = field_add(z[j], y_u[j]) - - - # Step 2: For surviving users, reconstruct b_u and subtract personal masks - for u in survived: - prg_seed_shares = [] - for reporter_id, reporter_shares in round4_data.items(): - target_key = str(u) - if target_key in reporter_shares and reporter_shares[target_key]['type'] == 'survived': - share_data = reporter_shares[target_key]['prg_seed_share'] - x = share_data[0] - y = bdecode(share_data[1]) - prg_seed_shares.append((x, y)) - - - if len(prg_seed_shares) >= THRESHOLD_CLIENTS: - b_u = Shamir.combine(prg_seed_shares[:THRESHOLD_CLIENTS]) - PRG_block = make_prg2(u, u, b_u, vec_len) - p_u = prg_block_to_field_elements(PRG_block, vec_len) - for j in range(vec_len): - z[j] = field_add(z[j], field_negate(p_u[j])) - - - # Step 3: For dropped users, reconstruct s_d^SK and subtract pairwise masks - for d in dropped: - s_sec_shares_1 = [] - s_sec_shares_2 = [] - for reporter_id, reporter_shares in round4_data.items(): - target_key = str(d) - if target_key in reporter_shares and reporter_shares[target_key]['type'] == 'dropped': - share_data = reporter_shares[target_key]['s_sec_share'] - s_sec_shares_1.append((share_data[0][0], bdecode(share_data[0][1]))) - s_sec_shares_2.append((share_data[1][0], bdecode(share_data[1][1]))) - - - if len(s_sec_shares_1) >= THRESHOLD_CLIENTS: - s_d_sk_half1 = Shamir.combine(s_sec_shares_1[:THRESHOLD_CLIENTS]) - s_d_sk_half2 = Shamir.combine(s_sec_shares_2[:THRESHOLD_CLIENTS]) - s_d_sk = X25519PrivateKey.from_private_bytes(s_d_sk_half1 + s_d_sk_half2) - - - # For each surviving user s, compute and subtract p_{s,d} - for s in survived: - s_s_pub = b64_to_pubkey(received_data[1][s]['key_s_pub']) - PRG_block = make_prg(d, s, s_d_sk, s_s_pub, vec_len) - prg_elements = prg_block_to_field_elements(PRG_block, vec_len) - - - for j in range(vec_len): - # p_{s,d}: what surviving user s added for dropped user d - if s > d: - p_s_d_j = prg_elements[j] - else: - p_s_d_j = field_negate(prg_elements[j]) - z[j] = field_add(z[j], field_negate(p_s_d_j)) - - - return z - +def extract_round_client_id_payload( + data: dict, expected_round: int +) -> tuple[int, int, dict]: + # Extract round number from JSON body + round = data.get("round") + # Handle errors + if round is None: + raise ValueError('Error: Missing "round" field in JSON body') + round = int(round) + assert ( + round == expected_round + ), f"Round number {expected_round} from POST does not match JSON body round number {round}" + if not (1 <= round <= ROUNDS): + raise ValueError(f"Error: Invalid round {round}; must be 1..{ROUNDS}") + + # Extract client_id from JSON body + client_id = data.get("client_id") + # TODO V2: HANDLE AUTHENTICATION! CURRENTLY WE DO NOT CHECK TO ENSURE CLIENT NOT LYING ABOUT ID. + # Handle errors + if client_id is None: + raise ValueError('Error: Missing "client_id" field in JSON body') + client_id = int(client_id) + if not (1 <= client_id <= MAX_CLIENTS): + raise ValueError( + f"Error: Invalid client_id {client_id}; must be 1..{MAX_CLIENTS}" + ) + # Extract payload from JSON body + payload = data.get("payload") + return (round, client_id, payload) + + +def build_keyset_response( + received_data: dict[int, dict[int, dict]], round1_responders: set[int] +) -> dict[int, dict]: + # received_data format: {round: {client_id: {payload}}} + round1_data = received_data[1] + return {client_id: round1_data[client_id] for client_id in round1_responders} + + +def build_sharekeys_response( + client_id: int, + received_data: dict[int, dict[int, dict]], + round2_responders: set[int], +) -> dict[int, dict]: + # received_data format: {round: {client_id: {payload}}} + # Received from client 3 in round 2: {'1': ['jNrWa3PZPHDvX391', 'xxx'], '2': ['kDUON98rzqdA98+n', 'xxx']} + # TODO: CHANGE RECEIVED DATA TO ONLY INPUTTING THIS ROUND'S DATA + round2_data = received_data[2] + to_ret = {} + for other_id in round2_responders: + if other_id == client_id: + continue + if other_id not in round2_data or str(client_id) not in round2_data[other_id]: + raise ValueError( + f"Error: build_sharekeys_response called for client {client_id}, but no round 2 data from other client {other_id} for this client" + ) + to_ret[other_id] = round2_data[other_id][ + str(client_id) + ] # add received nonce and ciphertext from other_id for client_id to client_id's dictionary + return to_ret + + +def build_masked_input_response( + client_id: int, + received_data: dict[int, dict[int, dict]], + round3_responders: set[int], +) -> dict: + return {"u2_users": sorted(list(round3_responders))} + + +def build_consistency_check_response( + received_data: dict, round3_responders: set[int], round4_responders: set[int] +) -> dict: + u3_list = sorted(list(round3_responders)) + signatures = {} + for uid in round4_responders: + signatures[str(uid)] = received_data[4][uid]["signature"] + return { + "users": {str(uid): {} for uid in round4_responders}, + "u2_users": u3_list, + "u2_signatures": signatures, + } + + +def compute_final_aggregate( + received_data, round2_responders, round3_responders, round5_data, vec_len +): + """ + Compute the final aggregate: z = sum(y_u) - sum(p_u) - sum(p_{s,d}) + where s are surviving users and d are dropped users. + """ + survived = set(round3_responders) + dropped = set(round2_responders) - survived + + # Step 1: Sum all y_u for u in U3 + z = [0] * vec_len + for u in round3_responders: + y_u = received_data[3][u] + for j in range(vec_len): + z[j] = field_add(z[j], y_u[j]) + + # Step 2: For surviving users, reconstruct b_u and subtract personal masks + for u in survived: + prg_seed_shares = [] + for reporter_id, reporter_shares in round5_data.items(): + target_key = str(u) + if ( + target_key in reporter_shares + and reporter_shares[target_key]["type"] == "survived" + ): + share_data = reporter_shares[target_key]["prg_seed_share"] + x = share_data[0] + y = bdecode(share_data[1]) + prg_seed_shares.append((x, y)) + + if len(prg_seed_shares) >= THRESHOLD_CLIENTS: + b_u = Shamir.combine(prg_seed_shares[:THRESHOLD_CLIENTS]) + PRG_block = make_prg2(u, u, b_u, vec_len) + p_u = prg_block_to_field_elements(PRG_block, vec_len) + for j in range(vec_len): + z[j] = field_add(z[j], field_negate(p_u[j])) + + # Step 3: For dropped users, reconstruct s_d^SK and subtract pairwise masks + for d in dropped: + s_sec_shares_1 = [] + s_sec_shares_2 = [] + for reporter_id, reporter_shares in round5_data.items(): + target_key = str(d) + if ( + target_key in reporter_shares + and reporter_shares[target_key]["type"] == "dropped" + ): + share_data = reporter_shares[target_key]["s_sec_share"] + s_sec_shares_1.append((share_data[0][0], bdecode(share_data[0][1]))) + s_sec_shares_2.append((share_data[1][0], bdecode(share_data[1][1]))) + + if len(s_sec_shares_1) >= THRESHOLD_CLIENTS: + s_d_sk_half1 = Shamir.combine(s_sec_shares_1[:THRESHOLD_CLIENTS]) + s_d_sk_half2 = Shamir.combine(s_sec_shares_2[:THRESHOLD_CLIENTS]) + s_d_sk = X25519PrivateKey.from_private_bytes(s_d_sk_half1 + s_d_sk_half2) + + # For each surviving user s, compute and subtract p_{s,d} + for s in survived: + s_s_pub = b64_to_pubkey(received_data[1][s]["key_s_pub"]) + PRG_block = make_prg(d, s, s_d_sk, s_s_pub, vec_len) + prg_elements = prg_block_to_field_elements(PRG_block, vec_len) + + for j in range(vec_len): + # p_{s,d}: what surviving user s added for dropped user d + if s > d: + p_s_d_j = prg_elements[j] + else: + p_s_d_j = field_negate(prg_elements[j]) + z[j] = field_add(z[j], field_negate(p_s_d_j)) + + return z def response_if_not_responder(client_id: int, round_failed_to_respond: int) -> dict: - return {'status': 'nonparticipant', 'message': f'Client {client_id} responded too late to participate in Round {round_failed_to_respond}.'} - + return { + "status": "nonparticipant", + "message": f"Client {client_id} responded too late to participate in Round {round_failed_to_respond}.", + } diff --git a/src/client.py b/src/client.py index d92eae7..0b445ce 100644 --- a/src/client.py +++ b/src/client.py @@ -8,339 +8,464 @@ from Crypto.Random import get_random_bytes from Crypto.Protocol.SecretSharing import Shamir from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey -from config import ROUNDS, THRESHOLD_CLIENTS, PRG_SEED_SIZE, DEBUG, DEBUG_TESTING_DELAY, FIELD_ELEMENT_SIZE, R -from _client_helper import bencode, bdecode, field_add, field_negate, jencode_to_bytes, jdecode_from_bytes, jencode_ciphertexts_for_other_r1r_r2, jdecode_ciphertexts, pubkey_to_b64, b64_to_pubkey, privkey_to_raw_bytes, do_round -from _client_helper import derive_shared_key, encrypt_with_derived_key, decrypt_with_derived_key, ids_to_associated_data, make_prg, make_prg2, prg_block_to_field_elements, field_negate, field_add - - -SERVER_URL = 'http://127.0.0.1:5000' +from cryptography.hazmat.primitives.asymmetric import ed25519 +from config import ( + ROUNDS, + THRESHOLD_CLIENTS, + PRG_SEED_SIZE, + DEBUG, + DEBUG_TESTING_DELAY, + FIELD_ELEMENT_SIZE, + R, +) +from _client_helper import ( + bencode, + bdecode, + field_add, + field_negate, + jencode_to_bytes, + jdecode_from_bytes, + jencode_ciphertexts_for_other_r1r_r2, + jdecode_ciphertexts, + pubkey_to_b64, + b64_to_pubkey, + privkey_to_raw_bytes, + do_round, +) +from _client_helper import ( + derive_shared_key, + encrypt_with_derived_key, + decrypt_with_derived_key, + ids_to_associated_data, + make_prg, + make_prg2, + prg_block_to_field_elements, + field_negate, + field_add, +) +from _client_helper import pubkey_to_bytes + +SERVER_URL = "http://127.0.0.1:5000" class SecureAggregationClient: - def __init__(self, client_id: int, x_u: list[int]): - self.client_id = client_id - self.x_u = x_u - self.key_c_pub = None - self.key_c_sec = None - self.key_s_pub = None - self.key_s_sec = None - self.round1_responders = list() # is a list because we need consistent indexing - self.pubkeys_for_r1r = {} - self.symmkeys_for_other_r1r = {} - self.s_sec_shares = None - self.prg_seed = None - self.prg_seed_shares = None - self.p_u_v = {} - self.p_u = None - self.round2_responders = list() - # Shares received from other users (decrypted in Round 3, used in Round 4) - self.received_s_sec_shares = {} # {sender_id: ((x1, y1), (x2, y2))} - self.received_prg_seed_shares = {} # {sender_id: (x, y)} - - - def advertise_keys(self) -> dict: - self.key_c_sec = X25519PrivateKey.generate() - self.key_c_pub = self.key_c_sec.public_key() - - - self.key_s_sec = X25519PrivateKey.generate() - self.key_s_pub = self.key_s_sec.public_key() - - - return { - 'client_id': self.client_id, - 'round': 1, - 'payload': { - 'key_c_pub': pubkey_to_b64(self.key_c_pub), - 'key_s_pub': pubkey_to_b64(self.key_s_pub) - } - } - - - def share_keys(self, r1_response: dict) -> dict: - if r1_response is None: - print(f"Client {self.client_id} did not receive any response from server for round 1, aborting share keys.", flush=True) - sys.exit(1) - - - # Set clients that responded in r1 and their keys - self.round1_responders = [int(r1r_str) for r1r_str in r1_response.keys()] - self.pubkeys_for_r1r = {int(r1r_str): r1_response[r1r_str] for r1r_str in r1_response.keys()} # dict with int keys instead of str - - - # Verify uniqueness of keys - all_c_pub = {self.pubkeys_for_r1r[r1r]["key_c_pub"] for r1r in self.round1_responders} - all_s_pub = {self.pubkeys_for_r1r[r1r]["key_s_pub"] for r1r in self.round1_responders} - # assert(len(set(all_c_pub).union(set(all_s_pub))) == 2*len(self.round1_responders)) - # assert(len(self.round1_responders) >= THRESHOLD_CLIENTS, f"Only got {len(self.round1_responders)} responders, needed {THRESHOLD_CLIENTS} for security, aborting") - - - # Generate random value - prg_seed = get_random_bytes(PRG_SEED_SIZE) # b_u in Bhowmick et al. 2017 - self.prg_seed = prg_seed - - - # Shamir sharings of prg_seed and s_sec - prg_seed_shares = Shamir.split(k=THRESHOLD_CLIENTS, - n=len(self.round1_responders), - secret=prg_seed) - - # Shamir shares formatted as ((1, 0xdeadbeef), (2, 0xdeadbeef), ...) - s_sec_shares_1 = Shamir.split(k=THRESHOLD_CLIENTS, - n=len(self.round1_responders), - secret=privkey_to_raw_bytes(self.key_s_sec)[:16]) - s_sec_shares_2 = Shamir.split(k=THRESHOLD_CLIENTS, - n=len(self.round1_responders), - secret=privkey_to_raw_bytes(self.key_s_sec)[16:]) - self.prg_seed_shares = prg_seed_shares - self.s_sec_shares = (s_sec_shares_1, s_sec_shares_2) - - - # Key agreement and build message to other users - ciphertexts_for_other_r1r_r2 = {} - for (r1r, i) in zip(self.round1_responders, range(len(self.round1_responders))): - if r1r == self.client_id: - continue - else: - # Derive and store shared key - derived_key = derive_shared_key(self.client_id, r1r, self.key_c_sec, - b64_to_pubkey(self.pubkeys_for_r1r[r1r]["key_c_pub"])) - self.symmkeys_for_other_r1r[r1r] = derived_key - - # Build messages for other users - - message_for_other_r1r_r2 = { - "s_sec_share_xs": (s_sec_shares_1[i][0], s_sec_shares_2[i][0]), - "s_sec_share_ys": (bencode(s_sec_shares_1[i][1]), - bencode(s_sec_shares_2[i][1])), - "prg_seed_share_x": prg_seed_shares[i][0], - "prg_seed_share_y": bencode(prg_seed_shares[i][1]) - } - print(f"message from {self.client_id} to {r1r}: {message_for_other_r1r_r2}", flush=True) - - associated_data = ids_to_associated_data(self.client_id, r1r) - - - ciphertexts_for_other_r1r_r2[r1r] = encrypt_with_derived_key( - derived_key, - jencode_to_bytes(message_for_other_r1r_r2), - associated_data - ) - print(f"ct from {self.client_id} to {r1r}: {ciphertexts_for_other_r1r_r2[r1r]}", flush=True) - - - print("jencoding: ", jencode_ciphertexts_for_other_r1r_r2(ciphertexts_for_other_r1r_r2), flush=True) - - r2_payload = { - 'client_id': self.client_id, - 'round': 2, - 'payload': jencode_ciphertexts_for_other_r1r_r2(ciphertexts_for_other_r1r_r2) - } - return r2_payload - - - def masked_input_collection(self, ciphertexts_from_server: dict) -> dict: - if ciphertexts_from_server is None: - print(f"Client {self.client_id} did not receive any ciphertexts from server for round 3, aborting masked input collection.", flush=True) - sys.exit(1) - - self.round2_responders = [int(r2r_str) for r2r_str in ciphertexts_from_server.keys()] - self.round2_responders.append(self.client_id) # add self to responders for this round - vec_len = len(self.x_u) - - - # Step 1: Decrypt ciphertexts from other users and store their shares - # (needed in Round 4 for unmasking) - for r2r in self.round2_responders: - if r2r == self.client_id: - continue - - - ct_data = ciphertexts_from_server.get(str(r2r), ciphertexts_from_server.get(r2r)) - nonce_b64, ciphertext_b64 = ct_data - nonce, ciphertext = jdecode_ciphertexts(nonce_b64, ciphertext_b64) - derived_key = self.symmkeys_for_other_r1r[r2r] - associated_data = ids_to_associated_data(r2r, self.client_id) - - - plaintext_bytes = decrypt_with_derived_key( - derived_key, - nonce, - ciphertext, - associated_data - ) - msg = jdecode_from_bytes(plaintext_bytes) - - - # Store s_sec shares (two halves) - self.received_s_sec_shares[r2r] = ( - (msg["s_sec_share_xs"][0], bdecode(msg["s_sec_share_ys"][0])), - (msg["s_sec_share_xs"][1], bdecode(msg["s_sec_share_ys"][1])) - ) - # Store prg_seed share - self.received_prg_seed_shares[r2r] = ( - msg["prg_seed_share_x"], - bdecode(msg["prg_seed_share_y"]) - ) - - - # Step 2: Compute pairwise masks p_{u,v} for all v in U2 \ {u} - for r2r in self.round2_responders: - if r2r == self.client_id: - continue - other_s_pub = b64_to_pubkey(self.pubkeys_for_r1r[r2r]["key_s_pub"]) - PRG_block = make_prg(self.client_id, r2r, self.key_s_sec, other_s_pub, vec_len) - prg_elements = prg_block_to_field_elements(PRG_block, vec_len) - p_u_v = [] - - - for elem in prg_elements: - if self.client_id > r2r: - p_u_v.append(elem) - else: - p_u_v.append(field_negate(elem)) - - - self.p_u_v[r2r] = p_u_v - - - # Step 3: Compute personal mask p_u = PRG(b_u) - PRG_block = make_prg2(self.client_id, self.client_id, self.prg_seed, vec_len) - self.p_u = prg_block_to_field_elements(PRG_block, vec_len) - - - # Step 4: Compute masked input y_u = x_u + p_u + sum_{v in U2\{u}} p_{u,v} (mod R) - y_u = [] - for j in range(vec_len): - val = self.x_u[j] % R - val = field_add(val, self.p_u[j]) - for v in self.round2_responders: - if v != self.client_id: - val = field_add(val, self.p_u_v[v][j]) - y_u.append(val) - - - r3_payload = { - 'client_id': self.client_id, - 'round': 3, - 'payload': y_u - } - return r3_payload - - def unmasking(self, users_from_server: dict) -> dict: - if users_from_server is None: - print(f"Client {self.client_id} did not receive any users from server for round 4, aborting unmasking.", flush=True) - sys.exit(1) - - self.round3_responders = [int(r3r_str) for r3r_str in users_from_server.keys()] - - #Determine who dropped (in U2 but not in U3) - dropped = set(self.round2_responders) - set(self.round3_responders) - survived = set(self.round3_responders) - - shares_to_send = {} - for v in dropped: - if v in self.received_s_sec_shares: - share1, share2 = self.received_s_sec_shares[v] - shares_to_send[str(v)] = { - 'type': 'dropped', - 's_sec_share':[ - [share1[0], bencode(share1[1])], - [share2[0], bencode(share2[1])] - ] - } - - for v in survived: - if v == self.client_id: - continue - if v in self.received_prg_seed_shares: - share = self.received_prg_seed_shares[v] - shares_to_send[str(v)] = { - 'type': 'survived', - 'prg_seed_share': [share[0], bencode(share[1])] - } - - r4_payload = { - 'client_id': self.client_id, - 'round': 4, - 'payload': shares_to_send - } - return r4_payload - + def __init__( + self, + client_id: int, + x_u: list[int], + isactive: bool = False, + signingkeyfile: str = None, + verificationkeysfile: str = None, + ): + self.client_id = client_id + self.x_u = x_u + self.isactive = isactive + self.key_c_pub = None + self.key_c_sec = None + self.key_s_pub = None + self.key_s_sec = None + self.round1_responders = list() # is a list because we need consistent indexing + self.pubkeys_for_r1r = {} + self.symmkeys_for_other_r1r = {} + self.s_sec_shares = None + self.prg_seed = None + self.prg_seed_shares = None + self.p_u_v = {} + self.p_u = None + self.round2_responders = list() + # Shares received from other users (decrypted in Round 3, used in Round 4) + self.received_s_sec_shares = {} # {sender_id: ((x1, y1), (x2, y2))} + self.received_prg_seed_shares = {} # {sender_id: (x, y)} + self.key_d_pub = None + self.key_d_sec = None + # get signing and verfiication keys for active adversary + if self.isactive: + if signingkeyfile is not None and verificationkeysfile is not None: + with open(signingkeyfile, "rb") as f: + self.signingkey = ed25519.Ed25519PrivateKey.from_private_bytes(f.read()) + self.verificationkeys = {} + # verfication keys: {"id": hex_public_key} + # Simply load using json.load + with open(verificationkeysfile, "rb") as f: + verification_keys = json.load(f) + for client_id, pubkey_hex in verification_keys.items(): + self.verificationkeys[int(client_id)] = ( + ed25519.Ed25519PublicKey.from_public_bytes( + bytes.fromhex(pubkey_hex) + ) + ) + + def advertise_keys(self) -> dict: + self.key_c_sec = X25519PrivateKey.generate() + self.key_c_pub = self.key_c_sec.public_key() + + self.key_s_sec = X25519PrivateKey.generate() + self.key_s_pub = self.key_s_sec.public_key() + + payload = { + "client_id": self.client_id, + "round": 1, + "payload": { + "key_c_pub": pubkey_to_b64(self.key_c_pub), + "key_s_pub": pubkey_to_b64(self.key_s_pub), + }, + } + # Generate a signature sigma_u = SIG.sign(d_u_sk, c_u_pk||s_u_pk) + if self.isactive: + message = pubkey_to_bytes(self.key_c_pub) + pubkey_to_bytes(self.key_s_pub) + signature = self.signingkey.sign(message) + payload["payload"]["signature"] = signature.hex() + + return payload + + def share_keys(self, r1_response: dict) -> dict: + if r1_response is None: + print( + f"Client {self.client_id} did not receive any response from server for round 1, aborting share keys.", + flush=True, + ) + sys.exit(1) + + # Set clients that responded in r1 and their keys + self.round1_responders = [int(r1r_str) for r1r_str in r1_response.keys()] + self.pubkeys_for_r1r = { + int(r1r_str): r1_response[r1r_str] for r1r_str in r1_response.keys() + } # dict with int keys instead of str + + # verify signatures of other users' keys + if self.isactive: + for r1r in self.round1_responders: + if r1r == self.client_id: + continue + r1r_data = self.pubkeys_for_r1r[r1r] + r1r_c_pub = b64_to_pubkey(r1r_data["key_c_pub"]) + r1r_s_pub = b64_to_pubkey(r1r_data["key_s_pub"]) + r1r_signature = bytes.fromhex(r1r_data["signature"]) + message = pubkey_to_bytes(r1r_c_pub) + pubkey_to_bytes(r1r_s_pub) + try: + self.verificationkeys[r1r].verify(r1r_signature, message) + except Exception as e: + print( + f"Client {self.client_id} failed to verify signature of client {r1r} in round 1, aborting. Error: {e}", + flush=True, + ) + sys.exit(1) + + # Verify uniqueness of keys + all_c_pub = { + self.pubkeys_for_r1r[r1r]["key_c_pub"] for r1r in self.round1_responders + } + all_s_pub = { + self.pubkeys_for_r1r[r1r]["key_s_pub"] for r1r in self.round1_responders + } + assert len(set(all_c_pub).union(set(all_s_pub))) == 2 * len( + self.round1_responders + ) + # assert(len(self.round1_responders) >= THRESHOLD_CLIENTS, f"Only got {len(self.round1_responders)} responders, needed {THRESHOLD_CLIENTS} for security, aborting") + + # Generate random value + prg_seed = get_random_bytes(PRG_SEED_SIZE) # b_u in Bhowmick et al. 2017 + self.prg_seed = prg_seed + + # Shamir sharings of prg_seed and s_sec + prg_seed_shares = Shamir.split( + k=THRESHOLD_CLIENTS, n=len(self.round1_responders), secret=prg_seed + ) + + # Shamir shares formatted as ((1, 0xdeadbeef), (2, 0xdeadbeef), ...) + s_sec_shares_1 = Shamir.split( + k=THRESHOLD_CLIENTS, + n=len(self.round1_responders), + secret=privkey_to_raw_bytes(self.key_s_sec)[:16], + ) + s_sec_shares_2 = Shamir.split( + k=THRESHOLD_CLIENTS, + n=len(self.round1_responders), + secret=privkey_to_raw_bytes(self.key_s_sec)[16:], + ) + self.prg_seed_shares = prg_seed_shares + self.s_sec_shares = (s_sec_shares_1, s_sec_shares_2) + + # Key agreement and build message to other users + ciphertexts_for_other_r1r_r2 = {} + for r1r, i in zip(self.round1_responders, range(len(self.round1_responders))): + if r1r == self.client_id: + continue + else: + # Derive and store shared key + derived_key = derive_shared_key( + self.client_id, + r1r, + self.key_c_sec, + b64_to_pubkey(self.pubkeys_for_r1r[r1r]["key_c_pub"]), + ) + self.symmkeys_for_other_r1r[r1r] = derived_key + + # Build messages for other users + + message_for_other_r1r_r2 = { + "s_sec_share_xs": (s_sec_shares_1[i][0], s_sec_shares_2[i][0]), + "s_sec_share_ys": ( + bencode(s_sec_shares_1[i][1]), + bencode(s_sec_shares_2[i][1]), + ), + "prg_seed_share_x": prg_seed_shares[i][0], + "prg_seed_share_y": bencode(prg_seed_shares[i][1]), + } + print( + f"message from {self.client_id} to {r1r}: {message_for_other_r1r_r2}", + flush=True, + ) + + associated_data = ids_to_associated_data(self.client_id, r1r) + + ciphertexts_for_other_r1r_r2[r1r] = encrypt_with_derived_key( + derived_key, + jencode_to_bytes(message_for_other_r1r_r2), + associated_data, + ) + print( + f"ct from {self.client_id} to {r1r}: {ciphertexts_for_other_r1r_r2[r1r]}", + flush=True, + ) + + print( + "jencoding: ", + jencode_ciphertexts_for_other_r1r_r2(ciphertexts_for_other_r1r_r2), + flush=True, + ) + + r2_payload = { + "client_id": self.client_id, + "round": 2, + "payload": jencode_ciphertexts_for_other_r1r_r2( + ciphertexts_for_other_r1r_r2 + ), + } + return r2_payload + + def masked_input_collection(self, ciphertexts_from_server: dict) -> dict: + if ciphertexts_from_server is None: + print( + f"Client {self.client_id} did not receive any ciphertexts from server for round 3, aborting masked input collection.", + flush=True, + ) + sys.exit(1) + + self.round2_responders = [ + int(r2r_str) for r2r_str in ciphertexts_from_server.keys() + ] + self.round2_responders.append( + self.client_id + ) # add self to responders for this round + vec_len = len(self.x_u) + + # Step 1: Decrypt ciphertexts from other users and store their shares + # (needed in Round 4 for unmasking) + for r2r in self.round2_responders: + if r2r == self.client_id: + continue + + ct_data = ciphertexts_from_server.get( + str(r2r), ciphertexts_from_server.get(r2r) + ) + nonce_b64, ciphertext_b64 = ct_data + nonce, ciphertext = jdecode_ciphertexts(nonce_b64, ciphertext_b64) + derived_key = self.symmkeys_for_other_r1r[r2r] + associated_data = ids_to_associated_data(r2r, self.client_id) + + plaintext_bytes = decrypt_with_derived_key( + derived_key, nonce, ciphertext, associated_data + ) + msg = jdecode_from_bytes(plaintext_bytes) + + # Store s_sec shares (two halves) + self.received_s_sec_shares[r2r] = ( + (msg["s_sec_share_xs"][0], bdecode(msg["s_sec_share_ys"][0])), + (msg["s_sec_share_xs"][1], bdecode(msg["s_sec_share_ys"][1])), + ) + # Store prg_seed share + self.received_prg_seed_shares[r2r] = ( + msg["prg_seed_share_x"], + bdecode(msg["prg_seed_share_y"]), + ) + + # Step 2: Compute pairwise masks p_{u,v} for all v in U2 \ {u} + for r2r in self.round2_responders: + if r2r == self.client_id: + continue + other_s_pub = b64_to_pubkey(self.pubkeys_for_r1r[r2r]["key_s_pub"]) + PRG_block = make_prg( + self.client_id, r2r, self.key_s_sec, other_s_pub, vec_len + ) + prg_elements = prg_block_to_field_elements(PRG_block, vec_len) + p_u_v = [] + + for elem in prg_elements: + if self.client_id > r2r: + p_u_v.append(elem) + else: + p_u_v.append(field_negate(elem)) + + self.p_u_v[r2r] = p_u_v + + # Step 3: Compute personal mask p_u = PRG(b_u) + PRG_block = make_prg2(self.client_id, self.client_id, self.prg_seed, vec_len) + self.p_u = prg_block_to_field_elements(PRG_block, vec_len) + + # Step 4: Compute masked input y_u = x_u + p_u + sum_{v in U2\{u}} p_{u,v} (mod R) + y_u = [] + for j in range(vec_len): + val = self.x_u[j] % R + val = field_add(val, self.p_u[j]) + for v in self.round2_responders: + if v != self.client_id: + val = field_add(val, self.p_u_v[v][j]) + y_u.append(val) + + r3_payload = {"client_id": self.client_id, "round": 3, "payload": y_u} + return r3_payload + + def consistency_check(self, r3_response: dict) -> dict: + u3_list = r3_response.get("u2_users", []) + if not u3_list or len(u3_list) < THRESHOLD_CLIENTS: + print(f"Consistency check: insufficient participants, aborting", flush=True) + sys.exit(1) + + self.u3_users = sorted(u3_list) + message_to_sign = json.dumps(self.u3_users).encode("utf-8") + payload = { + "client_id": self.client_id, + "round": 4, + "payload":{} + } + if self.isactive: + payload["payload"]["signature"] = self.signingkey.sign(message_to_sign).hex() + else: + payload["payload"]["signature"] = "dummy" + return payload + + def unmasking(self, r4_response: dict) -> dict: + users_from_server = r4_response.get("users", None) + if users_from_server is None: + print( + f"Client {self.client_id} did not receive any users from server for round 4, aborting unmasking.", + flush=True, + ) + sys.exit(1) + + round4_responders = [int(r4r_str) for r4r_str in users_from_server.keys()] + + if len(round4_responders) < THRESHOLD_CLIENTS: + print( + f"Secutity Alert: Only {len(round4_responders)} users survived, aborting.", + flush=True, + ) + sys.exit(1) + + # active adversary check: verify all signatures of the u3 list from round 3 + if self.isactive: + u3_list_from_server = r4_response.get("u2_users", []) + message_to_verify = json.dumps(u3_list_from_server).encode("utf-8") + for other_id, sig in r4_response["u2_signatures"].items(): + other_id_int = int(other_id) + if other_id_int == self.client_id: + continue + try: + self.verificationkeys[other_id_int].verify( + bytes.fromhex(sig), message_to_verify + ) + except Exception as e: + print( + f"Client {self.client_id} failed to verify signature of client {other_id} in round 4, aborting. Error: {e}", + flush=True, + ) + sys.exit(1) + + # Determine who dropped (in U2 but not in U3) + dropped = set(self.round2_responders) - set(self.u3_users) + survived = set(self.u3_users) + + shares_to_send = {} + for v in dropped: + if v in self.received_s_sec_shares: + share1, share2 = self.received_s_sec_shares[v] + shares_to_send[str(v)] = { + "type": "dropped", + "s_sec_share": [ + [share1[0], bencode(share1[1])], + [share2[0], bencode(share2[1])], + ], + } + + for v in survived: + if v == self.client_id: + continue + if v in self.received_prg_seed_shares: + share = self.received_prg_seed_shares[v] + shares_to_send[str(v)] = { + "type": "survived", + "prg_seed_share": [share[0], bencode(share[1])], + } + + r5_payload = { + "client_id": self.client_id, + "round": 5, + "payload": shares_to_send, + } + return r5_payload def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--id', type=int, required=True, help='Client id (1..i for testing)') - parser.add_argument('--vec', type=str, required="1,2,3", help='Input vector as comma-separated integers, e.g. "1,2,3"') - args = parser.parse_args() - client_id = args.id - - - x_u = [int(x) for x in args.vec.split(",")] - print(f"Client {client_id} starting with input vector {x_u}", flush=True) - - - # Setup: Initialize client - client = SecureAggregationClient(client_id=client_id, x_u=x_u) - - - # Round 1: Advertise Keys - r1_payload = client.advertise_keys() - testing_delay = (client_id == 1) if DEBUG_TESTING_DELAY else False # if we are delaying a client for testing, which one? - round_delay = 1 # if we are delaying a client for testing, which round to delay - r1_response = do_round(client_id, 1, r1_payload, SERVER_URL) - - # Abort if nonparticipant in Round 1 - if 'status' in r1_response and r1_response['status'] == 'nonparticipant': - print(f"Client {client_id} is a nonparticipant, aborting.", flush=True) - return - - #abort if client id is 4 or 5 - if client_id in [9, 10]: - print(f"Client {client_id} is dropping out.", flush=True) - return - - - # Round 2: Share Keys - r2_payload = client.share_keys(r1_response) - if DEBUG: print(f"Client {client_id} round 2 payload: {r2_payload}", flush=True) - r2_response = do_round(client_id, 2, r2_payload, SERVER_URL) - print(f"Client {client_id} round 2 response: {r2_response}", flush=True) - - #abort if client id is 4 or 5 - if client_id in [7, 8]: - print(f"Client {client_id} is dropping out.", flush=True) - return - - #Round 3: Masked Input Collection - r3_payload = client.masked_input_collection(r2_response) - print(f"Client {client_id} round 3 payload: {r3_payload}", flush=True) - r3_response = do_round(client_id, 3, r3_payload, SERVER_URL) - print(f"Client {client_id} round 3 response: {r3_response}", flush=True) - - #abort if client id is 4 or 5 - if client_id in [6]: - print(f"Client {client_id} is dropping out.", flush=True) - return - - #Round 4: Unmasking - r4_payload = client.unmasking(r3_response) - print(f"Client {client_id} round 4 payload: {r4_payload}", flush=True) - r4_response = do_round(client_id, 4, r4_payload, SERVER_URL) - print(f"Client {client_id} round 4 response: {r4_response}", flush=True) #this should print out the final answer - - # Remaining rounds -# for round in range(4, ROUNDS+1): -# payload = {'client_id': client_id, 'round': round, 'payload': f'Hello from client {client_id} round {round}'} -# print(f"Client {client_id}: Sending round {round}...", flush=True) -# resp = requests.post(f"{SERVER_URL}/round/{round}", json=payload) -# print(f"Client {client_id} round {round} response: {resp.json()}", flush=True) -# time.sleep(0.2) - - -if __name__ == '__main__': - main() - - - + parser = argparse.ArgumentParser() + parser.add_argument( + "--id", type=int, required=True, help="Client id (1..i for testing)" + ) + parser.add_argument( + "--vec", + type=str, + default="1,2,3", + help='Input vector as comma-separated integers, e.g. "1,2,3"', + ) + # each user rrecieve their signing key from the trusted third party + parser.add_argument( + "--signingkey", + type=str, + help="Path to file containing clients signing key for active adversary", + ) + # each user recieves the verification keys d_u_pk bound to each user iednitity v + parser.add_argument( + "--verificationkeys", + type=str, + help="Path to file containing dict of verification keys for all users for active adversary", + ) + + args = parser.parse_args() + client_id = args.id + + x_u = [int(x) for x in args.vec.split(",")] + print(f"Client {client_id} starting with input vector {x_u}", flush=True) + + # Setup: Initialize client + client = SecureAggregationClient( + client_id=client_id, + x_u=x_u, + isactive=True, + signingkeyfile=args.signingkey, + verificationkeysfile=args.verificationkeys, + ) + + r1_response = do_round(client_id, 1, client.advertise_keys(), SERVER_URL) + r2_response = do_round(client_id, 2, client.share_keys(r1_response), SERVER_URL) + r3_response = do_round( + client_id, 3, client.masked_input_collection(r2_response), SERVER_URL + ) + r4_response = do_round( + client_id, 4, client.consistency_check(r3_response), SERVER_URL + ) + r5_response = do_round(client_id, 5, client.unmasking(r4_response), SERVER_URL) + print(f"Aggregate: {r5_response.get('final_aggregate', None)}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/config.py b/src/config.py index d4d06f8..8b99e46 100644 --- a/src/config.py +++ b/src/config.py @@ -1,37 +1,48 @@ #!/usr/bin/env python3 + DEBUG = True DEBUG_TESTING_DELAY = False DEBUG_TESTING_DELAY_CLIENT_ID = 1 DEBUG_TESTING_DELAY_ROUND = 1 DEBUG_TESTING_DELAY_TIME = 6 -ROUNDS = 4 +ROUNDS = 5 -MAX_CLIENTS = 10 +MAX_CLIENTS = 10 # for demo only! Should really be more like 200 # n in Bonawitz et al. 2017 -THRESHOLD_CLIENTS = 3 + +THRESHOLD_CLIENTS = 3 # for demo only! Should really be more like 100 # t in Bonawitz et al. 2017 + # Timing parameters (general; overwritten by round-specific params as needed) -THRESHOLD_WAIT_ALL = 1 # Once threshold met, wait this many seconds before finalizing round i -POLL_INTERVAL_ALL = 0.5 # Once a client joins, poll every this many seconds to wait for round i result -MAX_POLLS_ALL = int(POLL_INTERVAL_ALL*20) # How many polls to wait maximum before giving up? TODO: For real world, should be MUCH higher, like 100+. +THRESHOLD_WAIT_ALL = ( + 1 # Once threshold met, wait this many seconds before finalizing round i +) +POLL_INTERVAL_ALL = ( + 0.5 # Once a client joins, poll every this many seconds to wait for round i result +) +MAX_POLLS_ALL = int( + POLL_INTERVAL_ALL * 20 +) # How many polls to wait maximum before giving up? TODO: For real world, should be MUCH higher, like 100+. + # Round-specific timing parameters -THRESHOLD_WAITS = {i: THRESHOLD_WAIT_ALL for i in range(1, ROUNDS+1)} -POLL_INTERVALS = {i: POLL_INTERVAL_ALL for i in range(1, ROUNDS+1)} -MAX_POLLS = {i: MAX_POLLS_ALL for i in range(1, ROUNDS+1)} +THRESHOLD_WAITS = {i: THRESHOLD_WAIT_ALL for i in range(1, ROUNDS + 1)} +POLL_INTERVALS = {i: POLL_INTERVAL_ALL for i in range(1, ROUNDS + 1)} +MAX_POLLS = {i: MAX_POLLS_ALL for i in range(1, ROUNDS + 1)} + -PRG_SEED_SIZE = 16 +PRG_SEED_SIZE = 16 # 128 bits # size of b_u in Bhowmick et al. 2017 + DERIVED_KEY_LENGTH = 32 - # We work on the field defined by prime p = 2^128 - 159, so field elements are 16 bytes long (128 bits) FIELD_ELEMENT_SIZE = 16 -R = 2**128 - 159 \ No newline at end of file +R = 2**128 - 159 diff --git a/src/server.py b/src/server.py index c25f55c..a49a2c0 100644 --- a/src/server.py +++ b/src/server.py @@ -3,225 +3,331 @@ import threading import time from config import ROUNDS, DEBUG, MAX_CLIENTS, THRESHOLD_CLIENTS, THRESHOLD_WAITS -from _server_helper import extract_round_client_id_payload, response_if_not_responder, build_keyset_response, build_sharekeys_response, build_masked_input_response +from _server_helper import ( + extract_round_client_id_payload, + response_if_not_responder, + build_keyset_response, + build_sharekeys_response, + build_masked_input_response, + build_consistency_check_response, +) from _server_helper import compute_final_aggregate + class SecureAggregationServer: - """Server for secure aggregation protocol.""" - - def __init__(self): - """Initialize the server with empty state and thread lock""" - self.received_data = {r: {} for r in range(1, ROUNDS+1)} # stores per-round received data - self.lock = threading.Lock() # locks state while edits being made - self.app = Flask(__name__) # Flask app instance to receive/send HTTP/POST requests (calls _setup_routes) - self.roundi_responders = {r: set() for r in range(1, ROUNDS+1)} # Clients that responded in each round - self.roundi_threshold_met = {r: threading.Event() for r in range(1, ROUNDS+1)} # Event to signal when threshold is reached - self.roundi_responders_locked = {r: threading.Event() for r in range(1, ROUNDS+1)} # Event to signal when threshold wait period ENDS - self.roundi_result = {r: None for r in range(1, ROUNDS+1)} # Computed result after each round's threshold wait - self.final_aggregate = None # Final aggregate result after round 4 - self._setup_routes() # Set up flask routes (called by Flask app setup) - - # TODO: Currently, round_num (sent as argument in POST) is redundant with "round" field sent in JSON body; we probably don't need both - def _setup_routes(self): - """Boilerplate to set up Flask routes to send/receive HTTP/POST""" - self.app.route('/round/', methods=['POST'])(self.handle_round) - self.app.route('/round1/result', methods=['GET'])(self.get_round1_result) - self.app.route('/round2/result', methods=['GET'])(self.get_round2_result) - self.app.route('/round3/result', methods=['GET'])(self.get_round3_result) - self.app.route('/round4/result', methods=['GET'])(self.get_round4_result) - - def handle_round(self, round_num: int): - """Handle incoming round data from clients""" - # Get data from the request - data = request.get_json(force=True) - - (round, client_id, payload) = extract_round_client_id_payload(data, expected_round=round_num) - - with self.lock: - self.received_data[round][client_id] = payload - - # Log received data from client - print(f"Received from client {client_id} in round {round}: {payload}", flush=True) - - # Run round-specific pre-checks - pre_check_response = None - if round == 1: - pre_check_response = self._check_round1(client_id) - elif round == 2: - pre_check_response = self._check_round2(client_id) - elif round == 3: - pre_check_response = self._check_round3(client_id) - elif round == 4: - pre_check_response = self._check_round4(client_id) - - # If pre-check failed, return the error response - if pre_check_response is not None: - return jsonify(pre_check_response) - - # Common threshold-based collection logic - response = self._collect_round_response(round, client_id, payload) - - # Send response - return jsonify(response) - - def _check_round1(self, client_id: int) -> dict: - """Pre-check for round 1. Returns None if checks pass, or error response if they fail.""" - return None # Round 1 has no special pre-checks - - def _check_round2(self, client_id: int) -> dict: - """Pre-check for round 2. Returns None if checks pass, or error response if they fail.""" - # Check if client participated in round 1 - if not self.roundi_responders_locked[1].is_set(): - return {'status': 'error', 'message': 'Round 2 received before round 1 threshold wait completed.'} - elif client_id not in self.roundi_responders[1]: - return response_if_not_responder(client_id, 1) - return None - - def _check_round3(self, client_id: int) -> dict: - """Pre-check for round 3. Returns None if checks pass, or error response if they fail.""" - # Check if client participated in round 2 - if not self.roundi_responders_locked[2].is_set(): - return {'status': 'error', 'message': 'Round 3 received before round 2 threshold wait completed.'} - elif client_id not in self.roundi_responders[2]: - return response_if_not_responder(client_id, 2) - return None - - def _check_round4(self, client_id: int) -> dict: - """Pre-check for round 4. Returns None if checks pass, or error response if they fail.""" - # Check if client participated in round 3 - if not self.roundi_responders_locked[3].is_set(): - return {'status': 'error', 'message': 'Round 4 received before round 3 threshold wait completed.'} - elif client_id not in self.roundi_responders[3]: - return response_if_not_responder(client_id, 3) - return None - - def _collect_round_response(self, round: int, client_id: int, payload) -> dict: - """Common logic for collecting clients in a round with threshold-based waiting.""" - max_responders = MAX_CLIENTS if round == 1 else len(self.roundi_responders[round-1]) - - # If we haven't hit the threshold yet, add this client to responders - if not self.roundi_responders_locked[round].is_set() and len(self.roundi_responders[round]) < max_responders: - self.roundi_responders[round].add(client_id) - print(f"Round {round}: Client {client_id} added. Responders so far: {len(self.roundi_responders[round])} / Threshold: {THRESHOLD_CLIENTS}", flush=True) - - # When we first hit the threshold, mark the event and start the wait period - if len(self.roundi_responders[round]) >= THRESHOLD_CLIENTS: - print(f"Round {round}: Threshold reached ({THRESHOLD_CLIENTS} clients). Starting wait period of {THRESHOLD_WAITS[round]} seconds.", flush=True) - if not self.roundi_threshold_met[round].is_set(): - self.roundi_threshold_met[round].set() - - # Start thread that will wait and lock responders - if not self.roundi_responders_locked[round].is_set(): - threshold_wait_thread = threading.Thread(target=self._threshold_wait, args=(round,)) - threshold_wait_thread.daemon = True - threshold_wait_thread.start() - - # Respond immediately; actual result happens in get_roundN_result - return {'status': 'ok', 'message': f'Client {client_id} registered. Waiting for at least {max(THRESHOLD_CLIENTS-len(self.roundi_responders[round]), 0)} more clients.'} - - elif self.roundi_responders_locked[round].is_set(): - # Responders locked, client is late - return response_if_not_responder(client_id, round) - else: - # Still collecting responses - return {'status': 'ok', 'message': f'Client {client_id} registered for round {round}.'} - - - def _threshold_wait(self, round: int): - """Wait for threshold wait period after threshold is reached, collecting any additional clients.""" - responders = self.roundi_responders[round] - lock_event = self.roundi_responders_locked[round] - - print(f"Starting thread waiting for round {round} threshold wait period ({THRESHOLD_WAITS[round]}s)...", flush=True) - time.sleep(THRESHOLD_WAITS[round]) - - with self.lock: - print(f"Round {round}: Wait period complete. Final responders: {responders}", flush=True) - - # Signal that the event is complete and lock responders - lock_event.set() - - def get_round1_result(self): - """Clients poll this endpoint to get the round 1 result once threshold wait completes.""" - # Extract client_id from query parameter - client_id = request.args.get('client_id', type=int) - - # Wait until round 1 is complete and result is ready - self.roundi_responders_locked[1].wait() - - with self.lock: - if client_id in self.roundi_responders[1]: - response = build_keyset_response(self.received_data, self.roundi_responders[1]) - else: - response = response_if_not_responder(client_id, 1) - return jsonify(response) - - def get_round2_result(self): - """Clients poll this endpoint to get the round 2 result once threshold wait completes.""" - print("Get round 2 result called", flush=True) - # Extract client_id from query parameter - client_id = request.args.get('client_id', type=int) - - # Wait until round 2 is complete and result is ready - self.roundi_responders_locked[2].wait() - - with self.lock: - if client_id not in self.roundi_responders[1]: - response = response_if_not_responder(client_id, 1) - elif client_id not in self.roundi_responders[2]: - response = response_if_not_responder(client_id, 2) - else: - response = build_sharekeys_response(client_id, self.received_data, self.roundi_responders[2]) - return jsonify(response) - - def get_round3_result(self): - """Clients poll this endpoint to get the round 3 result once threshold wait completes.""" - print("Get round 3 result called", flush=True) - # Extract client_id from query parameter - client_id = request.args.get('client_id', type=int) - - # Wait until round 3 is complete and result is ready - self.roundi_responders_locked[3].wait() - - with self.lock: - if client_id not in self.roundi_responders[1]: - response = response_if_not_responder(client_id, 1) - elif client_id not in self.roundi_responders[2]: - response = response_if_not_responder(client_id, 2) - elif client_id not in self.roundi_responders[3]: - response = response_if_not_responder(client_id, 3) - else: - response = build_masked_input_response(client_id, self.received_data, self.roundi_responders[3]) - return jsonify(response) - - def get_round4_result(self): - """Clients poll this endpoint to get the round 4 result once threshold wait completes.""" - print("Get round 4 result called", flush=True) - # Extract client_id from query parameter - client_id = request.args.get('client_id', type=int) - - # Wait until round 4 is complete and result is ready - self.roundi_responders_locked[4].wait() - - with self.lock: - if self.final_aggregate is None: - response = {'status': 'error', 'message': 'Final aggregate not ready yet.'} - some_y = next(iter(self.received_data[3].values())) - vec_len = len(some_y) - self.final_aggregate = compute_final_aggregate(self.received_data, self.roundi_responders[2], self.roundi_responders[3], self.received_data[4], vec_len) - if client_id not in self.roundi_responders[3]: - response = response_if_not_responder(client_id, 3) - elif client_id not in self.roundi_responders[4]: - response = response_if_not_responder(client_id, 4) - else: - response = {'status': 'ok', 'final_aggregate': self.final_aggregate} - return jsonify(response) - - def run(self, host='127.0.0.1', port=5000, debug=False): - """Start the Flask server""" - self.app.run(host=host, port=port, debug=debug) - - -if __name__ == '__main__': - server = SecureAggregationServer() - server.run(host='127.0.0.1', port=5000, debug=False) + """Server for secure aggregation protocol.""" + + def __init__(self): + """Initialize the server with empty state and thread lock""" + self.received_data = { + r: {} for r in range(1, ROUNDS + 1) + } # stores per-round received data + self.lock = threading.Lock() # locks state while edits being made + self.app = Flask( + __name__ + ) # Flask app instance to receive/send HTTP/POST requests (calls _setup_routes) + self.roundi_responders = { + r: set() for r in range(1, ROUNDS + 1) + } # Clients that responded in each round + self.roundi_threshold_met = { + r: threading.Event() for r in range(1, ROUNDS + 1) + } # Event to signal when threshold is reached + self.roundi_responders_locked = { + r: threading.Event() for r in range(1, ROUNDS + 1) + } # Event to signal when threshold wait period ENDS + self.roundi_result = { + r: None for r in range(1, ROUNDS + 1) + } # Computed result after each round's threshold wait + self.final_aggregate = None # Final aggregate result after round 4 + self._setup_routes() # Set up flask routes (called by Flask app setup) + + # TODO: Currently, round_num (sent as argument in POST) is redundant with "round" field sent in JSON body; we probably don't need both + def _setup_routes(self): + """Boilerplate to set up Flask routes to send/receive HTTP/POST""" + self.app.route("/round/", methods=["POST"])(self.handle_round) + self.app.route("/round1/result", methods=["GET"])(self.get_round1_result) + self.app.route("/round2/result", methods=["GET"])(self.get_round2_result) + self.app.route("/round3/result", methods=["GET"])(self.get_round3_result) + self.app.route("/round4/result", methods=["GET"])(self.get_round4_result) + self.app.route("/round5/result", methods=["GET"])(self.get_round5_result) + + def handle_round(self, round_num: int): + """Handle incoming round data from clients""" + # Get data from the request + data = request.get_json(force=True) + + (round, client_id, payload) = extract_round_client_id_payload( + data, expected_round=round_num + ) + + with self.lock: + self.received_data[round][client_id] = payload + + # Log received data from client + print( + f"Received from client {client_id} in round {round}: {payload}", + flush=True, + ) + + # Run round-specific pre-checks + pre_check_response = None + if round == 1: + pre_check_response = self._check_round1(client_id) + elif round == 2: + pre_check_response = self._check_round2(client_id) + elif round == 3: + pre_check_response = self._check_round3(client_id) + elif round == 4: + pre_check_response = self._check_round4(client_id) + elif round == 5: + pre_check_response = self._check_round5(client_id) + + # If pre-check failed, return the error response + if pre_check_response is not None: + return jsonify(pre_check_response) + + # Common threshold-based collection logic + response = self._collect_round_response(round, client_id, payload) + + # Send response + return jsonify(response) + + def _check_round1(self, client_id: int) -> dict: + """Pre-check for round 1. Returns None if checks pass, or error response if they fail.""" + return None # Round 1 has no special pre-checks + + def _check_round2(self, client_id: int) -> dict: + """Pre-check for round 2. Returns None if checks pass, or error response if they fail.""" + # Check if client participated in round 1 + if not self.roundi_responders_locked[1].is_set(): + return { + "status": "error", + "message": "Round 2 received before round 1 threshold wait completed.", + } + elif client_id not in self.roundi_responders[1]: + return response_if_not_responder(client_id, 1) + return None + + def _check_round3(self, client_id: int) -> dict: + """Pre-check for round 3. Returns None if checks pass, or error response if they fail.""" + # Check if client participated in round 2 + if not self.roundi_responders_locked[2].is_set(): + return { + "status": "error", + "message": "Round 3 received before round 2 threshold wait completed.", + } + elif client_id not in self.roundi_responders[2]: + return response_if_not_responder(client_id, 2) + return None + + def _check_round4(self, client_id: int) -> dict: + """Pre-check for round 4. Returns None if checks pass, or error response if they fail.""" + # Check if client participated in round 3 + if not self.roundi_responders_locked[3].is_set(): + return { + "status": "error", + "message": "Round 4 received before round 3 threshold wait completed.", + } + elif client_id not in self.roundi_responders[3]: + return response_if_not_responder(client_id, 3) + return None + + def _check_round5(self, client_id: int) -> dict: + """Pre-check for round 5. Returns None if checks pass, or error response if they fail.""" + # Check if client participated in round 4 + if not self.roundi_responders_locked[4].is_set(): + return { + "status": "error", + "message": "Round 5 received before round 4 threshold wait completed.", + } + elif client_id not in self.roundi_responders[4]: + return response_if_not_responder(client_id, 4) + return None + + def _collect_round_response(self, round: int, client_id: int, payload) -> dict: + """Common logic for collecting clients in a round with threshold-based waiting.""" + max_responders = ( + MAX_CLIENTS if round == 1 else len(self.roundi_responders[round - 1]) + ) + + # If we haven't hit the threshold yet, add this client to responders + if ( + not self.roundi_responders_locked[round].is_set() + and len(self.roundi_responders[round]) < max_responders + ): + self.roundi_responders[round].add(client_id) + print( + f"Round {round}: Client {client_id} added. Responders so far: {len(self.roundi_responders[round])} / Threshold: {THRESHOLD_CLIENTS}", + flush=True, + ) + + # When we first hit the threshold, mark the event and start the wait period + if len(self.roundi_responders[round]) >= THRESHOLD_CLIENTS: + print( + f"Round {round}: Threshold reached ({THRESHOLD_CLIENTS} clients). Starting wait period of {THRESHOLD_WAITS[round]} seconds.", + flush=True, + ) + if not self.roundi_threshold_met[round].is_set(): + self.roundi_threshold_met[round].set() + + # Start thread that will wait and lock responders + if not self.roundi_responders_locked[round].is_set(): + threshold_wait_thread = threading.Thread( + target=self._threshold_wait, args=(round,) + ) + threshold_wait_thread.daemon = True + threshold_wait_thread.start() + + # Respond immediately; actual result happens in get_roundN_result + return { + "status": "ok", + "message": f"Client {client_id} registered. Waiting for at least {max(THRESHOLD_CLIENTS-len(self.roundi_responders[round]), 0)} more clients.", + } + + elif self.roundi_responders_locked[round].is_set(): + # Responders locked, client is late + return response_if_not_responder(client_id, round) + else: + # Still collecting responses + return { + "status": "ok", + "message": f"Client {client_id} registered for round {round}.", + } + + def _threshold_wait(self, round: int): + """Wait for threshold wait period after threshold is reached, collecting any additional clients.""" + responders = self.roundi_responders[round] + lock_event = self.roundi_responders_locked[round] + + print( + f"Starting thread waiting for round {round} threshold wait period ({THRESHOLD_WAITS[round]}s)...", + flush=True, + ) + time.sleep(THRESHOLD_WAITS[round]) + + with self.lock: + print( + f"Round {round}: Wait period complete. Final responders: {responders}", + flush=True, + ) + + # Signal that the event is complete and lock responders + lock_event.set() + + def get_round1_result(self): + """Clients poll this endpoint to get the round 1 result once threshold wait completes.""" + # Extract client_id from query parameter + client_id = request.args.get("client_id", type=int) + + # Wait until round 1 is complete and result is ready + self.roundi_responders_locked[1].wait() + + with self.lock: + if client_id in self.roundi_responders[1]: + response = build_keyset_response( + self.received_data, self.roundi_responders[1] + ) + else: + response = response_if_not_responder(client_id, 1) + return jsonify(response) + + def get_round2_result(self): + """Clients poll this endpoint to get the round 2 result once threshold wait completes.""" + print("Get round 2 result called", flush=True) + # Extract client_id from query parameter + client_id = request.args.get("client_id", type=int) + + # Wait until round 2 is complete and result is ready + self.roundi_responders_locked[2].wait() + + with self.lock: + if client_id not in self.roundi_responders[1]: + response = response_if_not_responder(client_id, 1) + elif client_id not in self.roundi_responders[2]: + response = response_if_not_responder(client_id, 2) + else: + response = build_sharekeys_response( + client_id, self.received_data, self.roundi_responders[2] + ) + return jsonify(response) + + def get_round3_result(self): + """Clients poll this endpoint to get the round 3 result once threshold wait completes.""" + print("Get round 3 result called", flush=True) + # Extract client_id from query parameter + client_id = request.args.get("client_id", type=int) + + # Wait until round 3 is complete and result is ready + self.roundi_responders_locked[3].wait() + + with self.lock: + if client_id not in self.roundi_responders[1]: + response = response_if_not_responder(client_id, 1) + elif client_id not in self.roundi_responders[2]: + response = response_if_not_responder(client_id, 2) + elif client_id not in self.roundi_responders[3]: + response = response_if_not_responder(client_id, 3) + else: + response = build_masked_input_response( + client_id, self.received_data, self.roundi_responders[3] + ) + return jsonify(response) + + def get_round4_result(self): + """Clients poll this endpoint to get the round 4 (ConsistencyCheck) result once threshold wait completes.""" + print("Get round 4 result called", flush=True) + # Extract client_id from query parameter + client_id = request.args.get("client_id", type=int) + + # Wait until round 4 is complete and result is ready + self.roundi_responders_locked[4].wait() + + with self.lock: + if client_id not in self.roundi_responders[3]: + response = response_if_not_responder(client_id, 3) + elif client_id not in self.roundi_responders[4]: + response = response_if_not_responder(client_id, 4) + else: + response = build_consistency_check_response( + self.received_data, + self.roundi_responders[3], + self.roundi_responders[4], + ) + return jsonify(response) + + def get_round5_result(self): + """Clients poll this endpoint to get the round 5 (Unmasking) result with the final aggregate.""" + print("Get round 5 result called", flush=True) + # Extract client_id from query parameter + client_id = request.args.get("client_id", type=int) + + # Wait until round 5 is complete and result is ready + self.roundi_responders_locked[5].wait() + + with self.lock: + if self.final_aggregate is None: + some_y = next(iter(self.received_data[3].values())) + vec_len = len(some_y) + self.final_aggregate = compute_final_aggregate( + self.received_data, + self.roundi_responders[2], + self.roundi_responders[3], + self.received_data[5], + vec_len, + ) + if client_id not in self.roundi_responders[4]: + response = response_if_not_responder(client_id, 4) + elif client_id not in self.roundi_responders[5]: + response = response_if_not_responder(client_id, 5) + else: + response = {"status": "ok", "final_aggregate": self.final_aggregate} + return jsonify(response) + + def run(self, host="127.0.0.1", port=5000, debug=False): + """Start the Flask server""" + self.app.run(host=host, port=port, debug=debug) + + +if __name__ == "__main__": + server = SecureAggregationServer() + server.run(host="127.0.0.1", port=5000, debug=False) diff --git a/src/test_advertisekeys.py b/src/test_advertisekeys.py new file mode 100644 index 0000000..a9236f6 --- /dev/null +++ b/src/test_advertisekeys.py @@ -0,0 +1,136 @@ +import base64 +import json +import pytest +from cryptography.exceptions import InvalidSignature +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey +from cryptography.hazmat.primitives.asymmetric import ed25519 + +from _client_helper import b64_to_pubkey, pubkey_to_b64, pubkey_to_bytes +from client import SecureAggregationClient + +@pytest.fixture +def base_client(): + def _make_client(client_id=1, x_u=None, isactive=False, signing_key_file=None, verification_keys_file=None): + if x_u is None: + x_u = [0] * 10 # Default input vector of length 10 + return SecureAggregationClient( + client_id=client_id, + x_u=x_u or [0] * 10, + isactive=isactive, + signingkeyfile=signing_key_file, + verificationkeysfile=verification_keys_file, + ) + return _make_client + +@pytest.fixture +def active_client_context(tmp_path, base_client): + cid = 5 + sk = ed25519.Ed25519PrivateKey.generate() + + raw_sk = sk.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption()) + + vk_hex = sk.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw).hex() + + sign_path = tmp_path / "sign.bin" + verif_path = tmp_path / "verif.json" + print(f"Signing key path: {sign_path}") + print(f"Verification key path: {verif_path}") + sign_path.write_bytes(raw_sk) + verif_path.write_text(json.dumps({str(cid): vk_hex})) + + client = base_client(client_id=cid, isactive=True, signing_key_file=str(sign_path), verification_keys_file=str(verif_path)) + + return client, sk + +class TestPayloadStructure: + def test_top_level_keys(self, base_client): + p = base_client().advertise_keys() + assert set(p) == {"client_id", "round", "payload"} + + @pytest.mark.parametrize("cid", [0, 7, 2**31 - 1]) + def test_client_id_handling(self, base_client, cid): + p = base_client(client_id=cid).advertise_keys() + assert p["client_id"] == cid + + def test_round_is_1(self, base_client): + p = base_client().advertise_keys() + assert p["round"] == 1 + + def test_inner_payload_keys_passive(self, base_client): + p = base_client().advertise_keys()["payload"] + assert set(p) == {"key_c_pub", "key_s_pub"} + +class TestKeyGeneration: + def test_keys_are_valid_x25519(self, base_client): + p = base_client().advertise_keys()["payload"] + assert isinstance(b64_to_pubkey(p["key_c_pub"]), X25519PublicKey) + assert isinstance(b64_to_pubkey(p["key_s_pub"]), X25519PublicKey) + + def test_keys_are_32(self, base_client): + p = base_client().advertise_keys()["payload"] + for k in ("key_c_pub", "key_s_pub"): + assert len(base64.b64decode(p[k])) == 32 + + def test_instance_state_set(self, base_client): + client = base_client() + client.advertise_keys() + assert all(x is not None for x in [ + client.key_c_sec, client.key_c_pub, + client.key_s_sec, client.key_s_pub + ]) + + def test_payload_matches_instance_state(self, base_client): + client = base_client() + p = client.advertise_keys()["payload"] + assert p["key_c_pub"] == pubkey_to_b64(client.key_c_pub) + assert p["key_s_pub"] == pubkey_to_b64(client.key_s_pub) + + def test_repeated_calls_produce_fresh_keys(self, base_client): + client = base_client() + p1 = client.advertise_keys()["payload"] + p2 = client.advertise_keys()["payload"] + assert p1["key_c_pub"] != p2["key_c_pub"] + assert p1["key_s_pub"] != p2["key_s_pub"] + +class TestActiveSigning: + + def test_signature_present(self, active_client_context): + client, _ = active_client_context + assert "signature" in client.advertise_keys()["payload"] + + def test_signature_valid(self, active_client_context): + client, sk = active_client_context + p = client.advertise_keys() + inner = p["payload"] + msg = (pubkey_to_bytes(b64_to_pubkey(inner["key_c_pub"])) + pubkey_to_bytes(b64_to_pubkey(inner["key_s_pub"]))) + + sk.public_key().verify(bytes.fromhex(inner["signature"]), msg) + + def test_signatures_covers_both_keys(self, active_client_context): + client, sk = active_client_context + inner = client.advertise_keys()["payload"] + sig = bytes.fromhex(inner["signature"]) + with pytest.raises(InvalidSignature): + sk.public_key().verify(sig, b"\x00") + + def test_signature_changes_per_call(self, active_client_context): + client, sk = active_client_context + p1 = client.advertise_keys() + p2 = client.advertise_keys() + assert p1["payload"]["signature"] != p2["payload"]["signature"] + +class TestingEdgeCases: + def test_active_flag_without_key_files_raises_on_sign(self, base_client): + client = base_client(isactive=True, signing_key_file=None, verification_keys_file=None) + with pytest.raises(AttributeError): + client.advertise_keys() + + def test_empty_x_u(self, base_client): + p = base_client(x_u=[]).advertise_keys() + assert p["round"] == 1 diff --git a/src/test_client.py b/src/test_client.py new file mode 100644 index 0000000..f15d3de --- /dev/null +++ b/src/test_client.py @@ -0,0 +1,486 @@ +#!/usr/bin/env python3 +import base64 +import requests +import argparse +import time +import json +import sys +from Crypto.Random import get_random_bytes +from Crypto.Protocol.SecretSharing import Shamir +from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey +from cryptography.hazmat.primitives.asymmetric import ed25519 +from config import ( + ROUNDS, + THRESHOLD_CLIENTS, + PRG_SEED_SIZE, + DEBUG, + DEBUG_TESTING_DELAY, + FIELD_ELEMENT_SIZE, + R, +) +from _client_helper import ( + bencode, + bdecode, + field_add, + field_negate, + jencode_to_bytes, + jdecode_from_bytes, + jencode_ciphertexts_for_other_r1r_r2, + jdecode_ciphertexts, + pubkey_to_b64, + b64_to_pubkey, + privkey_to_raw_bytes, + do_round, +) +from _client_helper import ( + derive_shared_key, + encrypt_with_derived_key, + decrypt_with_derived_key, + ids_to_associated_data, + make_prg, + make_prg2, + prg_block_to_field_elements, + field_negate, + field_add, +) +from _client_helper import pubkey_to_bytes + +SERVER_URL = "http://127.0.0.1:5000" + + +class SecureAggregationClient: + def __init__( + self, + client_id: int, + x_u: list[int], + isactive: bool = False, + signingkeyfile: str = None, + verificationkeysfile: str = None, + ): + self.client_id = client_id + self.x_u = x_u + self.isactive = isactive + self.key_c_pub = None + self.key_c_sec = None + self.key_s_pub = None + self.key_s_sec = None + self.round1_responders = list() # is a list because we need consistent indexing + self.pubkeys_for_r1r = {} + self.symmkeys_for_other_r1r = {} + self.s_sec_shares = None + self.prg_seed = None + self.prg_seed_shares = None + self.p_u_v = {} + self.p_u = None + self.round2_responders = list() + # Shares received from other users (decrypted in Round 3, used in Round 4) + self.received_s_sec_shares = {} # {sender_id: ((x1, y1), (x2, y2))} + self.received_prg_seed_shares = {} # {sender_id: (x, y)} + self.key_d_pub = None + self.key_d_sec = None + # get signing and verfiication keys for active adversary + if self.isactive: + if signingkeyfile is not None and verificationkeysfile is not None: + with open(signingkeyfile, "rb") as f: + self.signingkey = ed25519.Ed25519PrivateKey.from_private_bytes(f.read()) + self.verificationkeys = {} + # verfication keys: {"id": hex_public_key} + # Simply load using json.load + with open(verificationkeysfile, "rb") as f: + verification_keys = json.load(f) + for client_id, pubkey_hex in verification_keys.items(): + self.verificationkeys[int(client_id)] = ( + ed25519.Ed25519PublicKey.from_public_bytes( + bytes.fromhex(pubkey_hex) + ) + ) + + def advertise_keys(self) -> dict: + self.key_c_sec = X25519PrivateKey.generate() + self.key_c_pub = self.key_c_sec.public_key() + + self.key_s_sec = X25519PrivateKey.generate() + self.key_s_pub = self.key_s_sec.public_key() + + payload = { + "client_id": self.client_id, + "round": 1, + "payload": { + "key_c_pub": pubkey_to_b64(self.key_c_pub), + "key_s_pub": pubkey_to_b64(self.key_s_pub), + }, + } + # Generate a signature sigma_u = SIG.sign(d_u_sk, c_u_pk||s_u_pk) + if self.isactive: + message = pubkey_to_bytes(self.key_c_pub) + pubkey_to_bytes(self.key_s_pub) + signature = self.signingkey.sign(message) + payload["payload"]["signature"] = signature.hex() + + return payload + + def share_keys(self, r1_response: dict) -> dict: + if r1_response is None: + print( + f"Client {self.client_id} did not receive any response from server for round 1, aborting share keys.", + flush=True, + ) + sys.exit(1) + + # Set clients that responded in r1 and their keys + self.round1_responders = [int(r1r_str) for r1r_str in r1_response.keys()] + self.pubkeys_for_r1r = { + int(r1r_str): r1_response[r1r_str] for r1r_str in r1_response.keys() + } # dict with int keys instead of str + + # verify signatures of other users' keys + if self.isactive: + for r1r in self.round1_responders: + if r1r == self.client_id: + continue + r1r_data = self.pubkeys_for_r1r[r1r] + r1r_c_pub = b64_to_pubkey(r1r_data["key_c_pub"]) + r1r_s_pub = b64_to_pubkey(r1r_data["key_s_pub"]) + r1r_signature = bytes.fromhex(r1r_data["signature"]) + message = pubkey_to_bytes(r1r_c_pub) + pubkey_to_bytes(r1r_s_pub) + try: + self.verificationkeys[r1r].verify(r1r_signature, message) + except Exception as e: + print( + f"Client {self.client_id} failed to verify signature of client {r1r} in round 1, aborting. Error: {e}", + flush=True, + ) + sys.exit(1) + + # Verify uniqueness of keys + all_c_pub = { + self.pubkeys_for_r1r[r1r]["key_c_pub"] for r1r in self.round1_responders + } + all_s_pub = { + self.pubkeys_for_r1r[r1r]["key_s_pub"] for r1r in self.round1_responders + } + assert len(set(all_c_pub).union(set(all_s_pub))) == 2 * len( + self.round1_responders + ) + # assert(len(self.round1_responders) >= THRESHOLD_CLIENTS, f"Only got {len(self.round1_responders)} responders, needed {THRESHOLD_CLIENTS} for security, aborting") + + # Generate random value + prg_seed = get_random_bytes(PRG_SEED_SIZE) # b_u in Bhowmick et al. 2017 + self.prg_seed = prg_seed + + # Shamir sharings of prg_seed and s_sec + prg_seed_shares = Shamir.split( + k=THRESHOLD_CLIENTS, n=len(self.round1_responders), secret=prg_seed + ) + + # Shamir shares formatted as ((1, 0xdeadbeef), (2, 0xdeadbeef), ...) + s_sec_shares_1 = Shamir.split( + k=THRESHOLD_CLIENTS, + n=len(self.round1_responders), + secret=privkey_to_raw_bytes(self.key_s_sec)[:16], + ) + s_sec_shares_2 = Shamir.split( + k=THRESHOLD_CLIENTS, + n=len(self.round1_responders), + secret=privkey_to_raw_bytes(self.key_s_sec)[16:], + ) + self.prg_seed_shares = prg_seed_shares + self.s_sec_shares = (s_sec_shares_1, s_sec_shares_2) + + # Key agreement and build message to other users + ciphertexts_for_other_r1r_r2 = {} + for r1r, i in zip(self.round1_responders, range(len(self.round1_responders))): + if r1r == self.client_id: + continue + else: + # Derive and store shared key + derived_key = derive_shared_key( + self.client_id, + r1r, + self.key_c_sec, + b64_to_pubkey(self.pubkeys_for_r1r[r1r]["key_c_pub"]), + ) + self.symmkeys_for_other_r1r[r1r] = derived_key + + # Build messages for other users + + message_for_other_r1r_r2 = { + "s_sec_share_xs": (s_sec_shares_1[i][0], s_sec_shares_2[i][0]), + "s_sec_share_ys": ( + bencode(s_sec_shares_1[i][1]), + bencode(s_sec_shares_2[i][1]), + ), + "prg_seed_share_x": prg_seed_shares[i][0], + "prg_seed_share_y": bencode(prg_seed_shares[i][1]), + } + print( + f"message from {self.client_id} to {r1r}: {message_for_other_r1r_r2}", + flush=True, + ) + + associated_data = ids_to_associated_data(self.client_id, r1r) + + ciphertexts_for_other_r1r_r2[r1r] = encrypt_with_derived_key( + derived_key, + jencode_to_bytes(message_for_other_r1r_r2), + associated_data, + ) + print( + f"ct from {self.client_id} to {r1r}: {ciphertexts_for_other_r1r_r2[r1r]}", + flush=True, + ) + + print( + "jencoding: ", + jencode_ciphertexts_for_other_r1r_r2(ciphertexts_for_other_r1r_r2), + flush=True, + ) + + r2_payload = { + "client_id": self.client_id, + "round": 2, + "payload": jencode_ciphertexts_for_other_r1r_r2( + ciphertexts_for_other_r1r_r2 + ), + } + return r2_payload + + def masked_input_collection(self, ciphertexts_from_server: dict) -> dict: + if ciphertexts_from_server is None: + print( + f"Client {self.client_id} did not receive any ciphertexts from server for round 3, aborting masked input collection.", + flush=True, + ) + sys.exit(1) + + self.round2_responders = [ + int(r2r_str) for r2r_str in ciphertexts_from_server.keys() + ] + self.round2_responders.append( + self.client_id + ) # add self to responders for this round + vec_len = len(self.x_u) + + # Step 1: Decrypt ciphertexts from other users and store their shares + # (needed in Round 4 for unmasking) + for r2r in self.round2_responders: + if r2r == self.client_id: + continue + + ct_data = ciphertexts_from_server.get( + str(r2r), ciphertexts_from_server.get(r2r) + ) + nonce_b64, ciphertext_b64 = ct_data + nonce, ciphertext = jdecode_ciphertexts(nonce_b64, ciphertext_b64) + derived_key = self.symmkeys_for_other_r1r[r2r] + associated_data = ids_to_associated_data(r2r, self.client_id) + + plaintext_bytes = decrypt_with_derived_key( + derived_key, nonce, ciphertext, associated_data + ) + msg = jdecode_from_bytes(plaintext_bytes) + + # Store s_sec shares (two halves) + self.received_s_sec_shares[r2r] = ( + (msg["s_sec_share_xs"][0], bdecode(msg["s_sec_share_ys"][0])), + (msg["s_sec_share_xs"][1], bdecode(msg["s_sec_share_ys"][1])), + ) + # Store prg_seed share + self.received_prg_seed_shares[r2r] = ( + msg["prg_seed_share_x"], + bdecode(msg["prg_seed_share_y"]), + ) + + # Step 2: Compute pairwise masks p_{u,v} for all v in U2 \ {u} + for r2r in self.round2_responders: + if r2r == self.client_id: + continue + other_s_pub = b64_to_pubkey(self.pubkeys_for_r1r[r2r]["key_s_pub"]) + PRG_block = make_prg( + self.client_id, r2r, self.key_s_sec, other_s_pub, vec_len + ) + prg_elements = prg_block_to_field_elements(PRG_block, vec_len) + p_u_v = [] + + for elem in prg_elements: + if self.client_id > r2r: + p_u_v.append(elem) + else: + p_u_v.append(field_negate(elem)) + + self.p_u_v[r2r] = p_u_v + + # Step 3: Compute personal mask p_u = PRG(b_u) + PRG_block = make_prg2(self.client_id, self.client_id, self.prg_seed, vec_len) + self.p_u = prg_block_to_field_elements(PRG_block, vec_len) + + # Step 4: Compute masked input y_u = x_u + p_u + sum_{v in U2\{u}} p_{u,v} (mod R) + y_u = [] + for j in range(vec_len): + val = self.x_u[j] % R + val = field_add(val, self.p_u[j]) + for v in self.round2_responders: + if v != self.client_id: + val = field_add(val, self.p_u_v[v][j]) + y_u.append(val) + + r3_payload = {"client_id": self.client_id, "round": 3, "payload": y_u} + return r3_payload + + def consistency_check(self, r3_response: dict) -> dict: + u3_list = r3_response.get("u2_users", []) + if not u3_list or len(u3_list) < THRESHOLD_CLIENTS: + print(f"Consistency check: insufficient participants, aborting", flush=True) + sys.exit(1) + + self.u3_users = sorted(u3_list) + message_to_sign = json.dumps(self.u3_users).encode("utf-8") + payload = { + "client_id": self.client_id, + "round": 4, + "payload":{} + } + if self.isactive: + payload["payload"]["signature"] = self.signingkey.sign(message_to_sign).hex() + else: + payload["payload"]["signature"] = "dummy" + return payload + + def unmasking(self, r4_response: dict) -> dict: + users_from_server = r4_response.get("users", None) + if users_from_server is None: + print( + f"Client {self.client_id} did not receive any users from server for round 4, aborting unmasking.", + flush=True, + ) + sys.exit(1) + + round4_responders = [int(r4r_str) for r4r_str in users_from_server.keys()] + + if len(round4_responders) < THRESHOLD_CLIENTS: + print( + f"Secutity Alert: Only {len(round4_responders)} users survived, aborting.", + flush=True, + ) + sys.exit(1) + + # active adversary check: verify all signatures of the u3 list from round 3 + if self.isactive: + u3_list_from_server = r4_response.get("u2_users", []) + message_to_verify = json.dumps(u3_list_from_server).encode("utf-8") + for other_id, sig in r4_response["u2_signatures"].items(): + other_id_int = int(other_id) + if other_id_int == self.client_id: + continue + try: + self.verificationkeys[other_id_int].verify( + bytes.fromhex(sig), message_to_verify + ) + except Exception as e: + print( + f"Client {self.client_id} failed to verify signature of client {other_id} in round 4, aborting. Error: {e}", + flush=True, + ) + sys.exit(1) + + # Determine who dropped (in U2 but not in U3) + dropped = set(self.round2_responders) - set(self.u3_users) + survived = set(self.u3_users) + + shares_to_send = {} + for v in dropped: + if v in self.received_s_sec_shares: + share1, share2 = self.received_s_sec_shares[v] + shares_to_send[str(v)] = { + "type": "dropped", + "s_sec_share": [ + [share1[0], bencode(share1[1])], + [share2[0], bencode(share2[1])], + ], + } + + for v in survived: + if v == self.client_id: + continue + if v in self.received_prg_seed_shares: + share = self.received_prg_seed_shares[v] + shares_to_send[str(v)] = { + "type": "survived", + "prg_seed_share": [share[0], bencode(share[1])], + } + + r5_payload = { + "client_id": self.client_id, + "round": 5, + "payload": shares_to_send, + } + return r5_payload + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--id", type=int, required=True, help="Client id (1..i for testing)" + ) + parser.add_argument( + "--vec", + type=str, + default="1,2,3", + help='Input vector as comma-separated integers, e.g. "1,2,3"', + ) + # each user rrecieve their signing key from the trusted third party + parser.add_argument( + "--signingkey", + type=str, + help="Path to file containing clients signing key for active adversary", + ) + # each user recieves the verification keys d_u_pk bound to each user iednitity v + parser.add_argument( + "--verificationkeys", + type=str, + help="Path to file containing dict of verification keys for all users for active adversary", + ) + + args = parser.parse_args() + client_id = args.id + + x_u = [int(x) for x in args.vec.split(",")] + print(f"Client {client_id} starting with input vector {x_u}", flush=True) + + # Setup: Initialize client + client = SecureAggregationClient( + client_id=client_id, + x_u=x_u, + isactive=True, + signingkeyfile=args.signingkey, + verificationkeysfile=args.verificationkeys, + ) + + r1_response = do_round(client_id, 1, client.advertise_keys(), SERVER_URL) + + if client_id in [9, 10]: + print(f"Client {client_id} is dropping out.", flush=True) + return + + r2_response = do_round(client_id, 2, client.share_keys(r1_response), SERVER_URL) + + if client_id in [7, 8]: + print(f"Client {client_id} is dropping out.", flush=True) + return + + r3_response = do_round( + client_id, 3, client.masked_input_collection(r2_response), SERVER_URL + ) + + if client_id in [6]: + print(f"Client {client_id} is dropping out.", flush=True) + return + + r4_response = do_round( + client_id, 4, client.consistency_check(r3_response), SERVER_URL + ) + r5_response = do_round(client_id, 5, client.unmasking(r4_response), SERVER_URL) + print(f"Aggregate: {r5_response.get('final_aggregate', None)}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/test_consistencycheck.py b/src/test_consistencycheck.py new file mode 100644 index 0000000..d54b088 --- /dev/null +++ b/src/test_consistencycheck.py @@ -0,0 +1,122 @@ +import json +from client import SecureAggregationClient +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import ed25519 + +# Assuming these exist in your project as per previous context +from _client_helper import b64_to_pubkey, pubkey_to_b64, pubkey_to_bytes + +@pytest.fixture +def base_client(): + def _make_client(client_id=1, x_u=None, isactive=False, signing_key_file=None, verification_keys_file=None): + if x_u is None: + x_u = [0] * 10 # Default input vector of length 10 + return SecureAggregationClient( + client_id=client_id, + x_u=x_u or [0] * 10, + isactive=isactive, + signingkeyfile=signing_key_file, + verificationkeysfile=verification_keys_file, + ) + return _make_client + +@pytest.fixture +def active_client_context(tmp_path, base_client): + cid = 5 + sk = ed25519.Ed25519PrivateKey.generate() + + raw_sk = sk.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption()) + + vk_hex = sk.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw).hex() + + cid_1 = 1 + sk_1= ed25519.Ed25519PrivateKey.generate() + + raw_sk_1 = sk_1.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption()) + + vk_hex_1 = sk_1.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw).hex() + + sign_path = tmp_path / "sign.bin" + verif_path = tmp_path / "verif.json" + print(f"Signing key path: {sign_path}") + print(f"Verification key path: {verif_path}") + sign_path.write_bytes(raw_sk) + verif_path.write_text(json.dumps({str(cid): vk_hex, str(cid_1): vk_hex_1})) + + client = base_client(client_id=cid, isactive=True, signing_key_file=str(sign_path), verification_keys_file=str(verif_path)) + + return client, sk, sk_1 + + +class TestConsistencyCheckPayload: + def test_top_level_keys(self, base_client): + # Mocking a valid r3_response to prevent exit + r3 = {"u2_users": list(range(10))} + p = base_client().consistency_check(r3) + assert set(p) == {"client_id", "round", "payload"} + assert p["round"] == 4 + + def test_payload_contains_signature(self, base_client): + r3 = {"u2_users": list(range(10))} + p = base_client().consistency_check(r3)["payload"] + assert "signature" in p + + def test_dummy_signature_for_passive_client(self, base_client): + r3 = {"u2_users": list(range(10))} + p = base_client(isactive=False).consistency_check(r3)["payload"] + assert p["signature"] == "dummy" + +class TestConsistencyLogic: + def test_aborts_on_insufficient_users(self, base_client): + # Testing the sys.exit(1) logic + # Note: THRESHOLD_CLIENTS must be defined in your client module + r3 = {"u2_users": [1, 2]} # Assuming threshold > 2 + with pytest.raises(SystemExit) as pytest_wrapped_e: + base_client().consistency_check(r3) + assert pytest_wrapped_e.type == SystemExit + assert pytest_wrapped_e.value.code == 1 + + def test_u3_list_is_sorted_in_state(self, base_client): + client = base_client() + unsorted_users = [10, 2, 5, 1] + client.consistency_check({"u2_users": unsorted_users}) + assert client.u3_users == [1, 2, 5, 10] + +class TestActiveConsistencySigning: + def test_signature_valid_for_active_client(self, active_client_context): + # Correctly unpack the 3-tuple returned by the fixture + client, sk, _ = active_client_context + + # Ensure THRESHOLD_CLIENTS is small enough for this list + u3_list = [1, 5, 10, 12, 15] + + p = client.consistency_check({"u2_users": u3_list}) + sig_hex = p["payload"]["signature"] + + expected_message = json.dumps(sorted(u3_list)).encode("utf-8") + + # This will now work because sk is correctly assigned + sk.public_key().verify(bytes.fromhex(sig_hex), expected_message) + + def test_signature_fails_on_tampered_list(self, active_client_context): + client, sk, _ = active_client_context + u3_list = [1, 2, 3, 4, 5, 6] + + p = client.consistency_check({"u2_users": u3_list}) + sig_hex = p["payload"]["signature"] + + tampered_message = json.dumps([1, 2, 3, 4, 5, 7]).encode("utf-8") + + with pytest.raises(Exception): + sk.public_key().verify(bytes.fromhex(sig_hex), tampered_message) \ No newline at end of file diff --git a/src/test_maskedinputcollection.py b/src/test_maskedinputcollection.py new file mode 100644 index 0000000..e55268a --- /dev/null +++ b/src/test_maskedinputcollection.py @@ -0,0 +1,128 @@ +import pytest +import json +import base64 +from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey +from cryptography.hazmat.primitives import serialization + +# Assuming these helpers exist in your project +from _client_helper import ( + pubkey_to_b64, + encrypt_with_derived_key, + ids_to_associated_data, + jencode_to_bytes, +) +from client import SecureAggregationClient + +# Constants for field math (Adjust based on your actual config) +R = 2**32 # Example field size + +@pytest.fixture +def base_client(): + def _make_client(client_id=1, x_u=None, isactive=False, signing_key_file=None, verification_keys_file=None): + if x_u is None: + x_u = [0] * 10 # Default input vector of length 10 + return SecureAggregationClient( + client_id=client_id, + x_u=x_u or [0] * 10, + isactive=isactive, + signingkeyfile=signing_key_file, + verificationkeysfile=verification_keys_file, + ) + return _make_client + +@pytest.fixture +def client_with_state(base_client): + """ + Initializes a client and manually populates the internal state + required to enter Round 3. + """ + client = base_client(client_id=10, x_u=[100, 200, 300]) + + # 1. Setup own keys/seed + client.key_s_sec = X25519PrivateKey.generate() + client.key_s_pub = client.key_s_sec.public_key() + client.prg_seed = b"a_very_secret_seed_32_bytes_long" + + # 2. Setup knowledge of other clients (Round 1 leftovers) + # We'll simulate one other client (ID: 20) + other_sk = X25519PrivateKey.generate() + client.pubkeys_for_r1r = { + 20: {"key_s_pub": pubkey_to_b64(other_sk.public_key())} + } + + # 3. Setup symmetric keys for decryption (Round 2 leftovers) + client.symmkeys_for_other_r1r = { + 20: b"symmetric_key_between_10_and_20_" + } + + return client, other_sk + +class TestMaskedInputStructure: + def test_top_level_keys(self, client_with_state): + client, _ = client_with_state + # Empty dict or mock dict needed for input + res = client.masked_input_collection({}) + assert set(res.keys()) == {"client_id", "round", "payload"} + assert res["round"] == 3 + + def test_payload_is_list(self, client_with_state): + client, _ = client_with_state + res = client.masked_input_collection({}) + assert isinstance(res["payload"], list) + assert len(res["payload"]) == len(client.x_u) + +class TestMaskingLogic: + def test_decryption_and_storage(self, client_with_state): + client, other_sk = client_with_state + + # Create a mock ciphertext from client 20 to client 10 + inner_msg = { + "s_sec_share_xs": [1, 2], + "s_sec_share_ys": [base64.b64encode(b"share1").decode(), base64.b64encode(b"share2").decode()], + "prg_seed_share_x": 3, + "prg_seed_share_y": base64.b64encode(b"share3").decode() + } + + ad = ids_to_associated_data(20, client.client_id) + nonce, ct = encrypt_with_derived_key( + client.symmkeys_for_other_r1r[20], + jencode_to_bytes(inner_msg), + ad + ) + + server_input = { + "20": [base64.b64encode(nonce).decode(), base64.b64encode(ct).decode()] + } + + client.masked_input_collection(server_input) + + # Verify shares were decrypted and stored correctly + assert 20 in client.received_s_sec_shares + assert client.received_s_sec_shares[20][0][1] == b"share1" + + def test_masking_math_consistency(self, client_with_state): + """ + Verify that calling the function twice with same state results + in the same masked output (determinism). + """ + client, _ = client_with_state + res1 = client.masked_input_collection({})["payload"] + + # Reset ephemeral state that gets modified during the call if necessary + # Note: In your provided code, p_u_v and p_u are overwritten, so it's fine. + res2 = client.masked_input_collection({})["payload"] + + assert res1 == res2 + +class TestEdgeCases: + def test_null_server_input_aborts(self, client_with_state): + client, _ = client_with_state + with pytest.raises(SystemExit): + client.masked_input_collection(None) + + def test_empty_responders(self, client_with_state): + client, _ = client_with_state + # Even with no other responders, it should still apply the personal mask (p_u) + res = client.masked_input_collection({}) + # y_u = x_u + p_u + assert len(res["payload"]) == 3 \ No newline at end of file diff --git a/src/test_sharekeys.py b/src/test_sharekeys.py new file mode 100644 index 0000000..4143ff6 --- /dev/null +++ b/src/test_sharekeys.py @@ -0,0 +1,150 @@ +import pytest +import json +import base64 +from cryptography.hazmat.primitives.asymmetric import x25519, ed25519 +from cryptography.hazmat.primitives import serialization +from _client_helper import pubkey_to_b64, pubkey_to_bytes +from client import SecureAggregationClient + +@pytest.fixture +def base_client(): + def _make_client(client_id=1, x_u=None, isactive=False, signing_key_file=None, verification_keys_file=None): + if x_u is None: + x_u = [0] * 10 # Default input vector of length 10 + return SecureAggregationClient( + client_id=client_id, + x_u=x_u or [0] * 10, + isactive=isactive, + signingkeyfile=signing_key_file, + verificationkeysfile=verification_keys_file, + ) + return _make_client + +@pytest.fixture +def active_client_context(tmp_path, base_client): + cid = 5 + sk = ed25519.Ed25519PrivateKey.generate() + + raw_sk = sk.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption()) + + vk_hex = sk.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw).hex() + + cid_1 = 1 + sk_1= ed25519.Ed25519PrivateKey.generate() + + raw_sk_1 = sk_1.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption()) + + vk_hex_1 = sk_1.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw).hex() + + sign_path = tmp_path / "sign.bin" + verif_path = tmp_path / "verif.json" + print(f"Signing key path: {sign_path}") + print(f"Verification key path: {verif_path}") + sign_path.write_bytes(raw_sk) + verif_path.write_text(json.dumps({str(cid): vk_hex, str(cid_1): vk_hex_1})) + + client = base_client(client_id=cid, isactive=True, signing_key_file=str(sign_path), verification_keys_file=str(verif_path)) + + return client, sk, sk_1 + + +@pytest.fixture +def mock_r1_response(): + """Generates a dictionary representing keys from 3 clients.""" + response = {} + client_keys = {} # Store keys to verify signatures later + + for cid in [1, 2, 5]: # Client 5 is often our 'local' client + c_priv = x25519.X25519PrivateKey.generate() + s_priv = x25519.X25519PrivateKey.generate() + c_pub = c_priv.public_key() + s_pub = s_priv.public_key() + + response[str(cid)] = { + "key_c_pub": pubkey_to_b64(c_pub), + "key_s_pub": pubkey_to_b64(s_pub) + } + return response + +class TestShareKeysInputs: + def test_abort_on_none_response(self, base_client): + client = base_client() + with pytest.raises(SystemExit): + client.share_keys(None) + + def test_responder_list_initialization(self, base_client, mock_r1_response): + client = base_client(client_id=5) + # We need to run advertise_keys first to initialize local keys + client.advertise_keys() + + client.share_keys(mock_r1_response) + assert set(client.round1_responders) == {1, 2, 5} + assert 1 in client.pubkeys_for_r1r + +class TestShareKeysVerification: + def test_fails_on_invalid_signature(self, active_client_context, mock_r1_response): + client, _, _ = active_client_context + client.advertise_keys() + + # Add a signature field to one of the peer's data, but make it garbage + mock_r1_response["1"]["signature"] = "deadbeef" * 8 + + with pytest.raises(SystemExit): + client.share_keys(mock_r1_response) + + def test_passes_with_valid_signatures(self, active_client_context): + client, sk, sk1 = active_client_context + client.advertise_keys() + + # Setup a response where client 1 has a valid signature + c_pub = x25519.X25519PublicKey.from_public_bytes(b"\x01"*32) + s_pub = x25519.X25519PublicKey.from_public_bytes(b"\x02"*32) + msg = pubkey_to_bytes(c_pub) + pubkey_to_bytes(s_pub) + sig = sk1.sign(msg).hex() + + valid_response = { + "1": { + "key_c_pub": pubkey_to_b64(c_pub), + "key_s_pub": pubkey_to_b64(s_pub), + "signature": sig + }, + "5": { # Local client + "key_c_pub": pubkey_to_b64(client.key_c_pub), + "key_s_pub": pubkey_to_b64(client.key_s_pub), + "signature": "doesn't matter for self" + } + } + + # This should not raise SystemExit + client.share_keys(valid_response) + +class TestSecretSharing: + def test_uniqueness_constraint(self, base_client, mock_r1_response): + client = base_client(client_id=5) + client.advertise_keys() + + # Force a duplicate key to trigger the assert + dup_key = mock_r1_response["1"]["key_c_pub"] + mock_r1_response["2"]["key_c_pub"] = dup_key + + with pytest.raises(AssertionError): + client.share_keys(mock_r1_response) + + def test_prg_seed_generation(self, base_client, mock_r1_response): + client = base_client(client_id=5) + client.advertise_keys() + + client.share_keys(mock_r1_response) + assert hasattr(client, 'prg_seed') + assert len(client.prg_seed) > 0 + diff --git a/src/test_unmasking.py b/src/test_unmasking.py new file mode 100644 index 0000000..e82fb99 --- /dev/null +++ b/src/test_unmasking.py @@ -0,0 +1,174 @@ +import pytest +import json +import sys +from client import SecureAggregationClient +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import ed25519 +from unittest.mock import MagicMock + +# Assuming these exist in your helper or are accessible +# from _client_helper import bencode + +@pytest.fixture +def base_client(): + def _make_client(client_id=1, x_u=None, isactive=False, signing_key_file=None, verification_keys_file=None): + if x_u is None: + x_u = [0] * 10 # Default input vector of length 10 + return SecureAggregationClient( + client_id=client_id, + x_u=x_u or [0] * 10, + isactive=isactive, + signingkeyfile=signing_key_file, + verificationkeysfile=verification_keys_file, + ) + return _make_client + +@pytest.fixture +def active_client_context(tmp_path, base_client): + cid = 5 + sk = ed25519.Ed25519PrivateKey.generate() + + raw_sk = sk.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption()) + + vk_hex = sk.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw).hex() + + cid_1 = 1 + sk_1= ed25519.Ed25519PrivateKey.generate() + + raw_sk_1 = sk_1.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption()) + + vk_hex_1 = sk_1.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw).hex() + + sign_path = tmp_path / "sign.bin" + verif_path = tmp_path / "verif.json" + print(f"Signing key path: {sign_path}") + print(f"Verification key path: {verif_path}") + sign_path.write_bytes(raw_sk) + verif_path.write_text(json.dumps({str(cid): vk_hex, str(cid_1): vk_hex_1})) + + client = base_client(client_id=cid, isactive=True, signing_key_file=str(sign_path), verification_keys_file=str(verif_path)) + + return client, sk, sk_1 + + +@pytest.fixture +def client_with_state(active_client_context, monkeypatch): + import client as client_mod + # Set threshold low so tests don't exit + monkeypatch.setattr(client_mod, "THRESHOLD_CLIENTS", 1) + + # Unpack all 3 values from the previous fixture + client, sk, sk_1 = active_client_context + + client.round2_responders = [1, 2, 3, 5] + client.u3_users = [2, 5] + + # Mock shares (ensure these match the IDs in round2/u3) + client.received_s_sec_shares = { + 1: [(1, b"share1_a"), (2, b"share1_b")], + 3: [(1, b"share3_a"), (2, b"share3_b")] + } + client.received_prg_seed_shares = { + 2: (1, b"seed_share2") + } + + # Ensure verification keys exist for active check + client.verificationkeys = { + 2: MagicMock(), + 1: MagicMock() + } + return client + +class TestUnmaskingPayloadStructure: + def test_top_level_keys(self, client_with_state): + r4_response = { + "users": {"2": {}, "5": {}}, + "u2_users": [2, 5], + "u2_signatures": {} + } + p = client_with_state.unmasking(r4_response) + assert set(p) == {"client_id", "round", "payload"} + assert p["round"] == 5 + + def test_empty_users_aborts(self, client_with_state): + with pytest.raises(SystemExit): + client_with_state.unmasking({"users": None}) + +class TestUnmaskingLogic: + def test_shares_categorization(self, client_with_state): + """Verify dropped users get s_sec_shares and survivors get prg_seed_shares.""" + r4_response = { + "users": {"2": {}, "5": {}}, # Survived list from server + "u2_users": [2, 5], + "u2_signatures": {} + } + + # We need to ensure THRESHOLD_CLIENTS isn't hit. + # For testing, you might need to monkeypatch the constant: + # import client as client_mod + # client_mod.THRESHOLD_CLIENTS = 1 + + p = client_with_state.unmasking(r4_response)["payload"] + + # User 1 and 3 were in round2 but not u3 -> dropped + assert p["1"]["type"] == "dropped" + assert "s_sec_share" in p["1"] + + # User 2 was in u3 -> survived + assert p["2"]["type"] == "survived" + assert "prg_seed_share" in p["2"] + + # Client should not send shares for themselves + assert "5" not in p + +class TestUnmaskingActiveSecurity: + def test_signature_verification_called(self, client_with_state): + client_with_state.isactive = True + u2_users = [2, 5] + # Create a dummy signature + sig_hex = "00" * 64 + + r4_response = { + "users": {"2": {}, "5": {}}, + "u2_users": u2_users, + "u2_signatures": {"2": sig_hex} + } + + client_with_state.unmasking(r4_response) + + # Check if verify was called on the public key of user 2 + client_with_state.verificationkeys[2].verify.assert_called_once() + + def test_invalid_signature_aborts(self, client_with_state): + client_with_state.isactive = True + # Setup mock to raise error + client_with_state.verificationkeys[2].verify.side_effect = Exception("Invalid") + + r4_response = { + "users": {"2": {}, "5": {}}, + "u2_users": [2, 5], + "u2_signatures": {"2": "bad_sig"} + } + + with pytest.raises(SystemExit): + client_with_state.unmasking(r4_response) + +class TestUnmaskingEdgeCases: + def test_threshold_failure(self, client_with_state, monkeypatch): + # Force threshold to be high + import client + monkeypatch.setattr(client, "THRESHOLD_CLIENTS", 100) + + r4_response = {"users": {"2": {}, "5": {}}} + with pytest.raises(SystemExit): + client_with_state.unmasking(r4_response) \ No newline at end of file diff --git a/src/ttp.py b/src/ttp.py new file mode 100644 index 0000000..3b27c83 --- /dev/null +++ b/src/ttp.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +import argparse +import json +import os +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import ed25519 + +def generate_keys(n): + output_dir = "keys" + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + all_verification_keys = {} + for i in range(1, n+1): + private_key = ed25519.Ed25519PrivateKey.generate() + public_key = private_key.public_key() + + priv_bytes = private_key.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption() + ) + + with open(os.path.join(output_dir, f"sign-{i}.key"), "wb") as f: + f.write(priv_bytes) + + pub_bytes = public_key.public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw + ) + + all_verification_keys[i] = pub_bytes.hex() + + print(f"Generated keys for client {i}") + + with open(os.path.join(output_dir, "verification_keys.json"), "w") as f: + json.dump(all_verification_keys, f, indent=4) + + print(f"keys saved in '{output_dir}/'") + print("Distribute one 'sign-X.key' to each client and 'allverify.key' to everyone.") + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Offline PKI Key Generator for SecAgg.") + parser.add_argument('--N', type=int, required=True, help='Number of clients to generate keys for') + args = parser.parse_args() + generate_keys(args.N) \ No newline at end of file diff --git a/test_run_all.sh b/test_run_all.sh new file mode 100755 index 0000000..ae0402c --- /dev/null +++ b/test_run_all.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +# Run Flask server and three clients, collecting logs. +set -euo pipefail + +# Set ROOT_DIR to absolute path of directory containing this script +# Set SRC_DIR to ROOT_DIR/src +ROOT_DIR="$(cd "$(dirname "$0")" && pwd)" +cd "$ROOT_DIR" +SRC_DIR="$ROOT_DIR/src" +LOG_DIR="$ROOT_DIR/log" + +mkdir -p "$LOG_DIR" + +PYTHON_BIN="$SRC_DIR/.venv/bin/python" +if [[ ! -x "$PYTHON_BIN" ]]; then + PYTHON_BIN="python" +fi + + +# Start server +echo "Starting server, logging to $LOG_DIR/server.log..." +"$PYTHON_BIN" "$SRC_DIR/server.py" > "$LOG_DIR/server.log" 2>&1 & +SERVER_PID=$! + + +# Wait +sleep 1 + + +# Generate signing/verification keys +echo "Generating signing/verification keys..." +"$PYTHON_BIN" "$SRC_DIR/ttp.py" --N 10 + + +# Start clients +CLIENT_PIDS="" +for i in 1 2 3 4 5 6 7 8 9 10; do + echo "Starting client $i, logging to $LOG_DIR/client${i}.log..." + "$PYTHON_BIN" "$SRC_DIR/test_client.py" --id "$i" --vec "1,2,4" --signingkey "keys/sign-${i}.key" --verificationkeys "keys/verification_keys.json" > "$LOG_DIR/client${i}.log" 2>&1 & + CLIENT_PIDS="$CLIENT_PIDS $!" + sleep 0.1 +done + + +# Wait for clients to finish +for client_pid in $CLIENT_PIDS; do + wait "$client_pid" || true + echo "Client $client_pid finished." + sleep 0.1 +done + + +# Stop server +kill "$SERVER_PID" || true +echo "Server $SERVER_PID stopped." + + +# Done +echo "Done." + + + +