Skip to content

Commit cbd6578

Browse files
mahbdmahbd
committed
added zeuz secret sharing code
1 parent 0caa12d commit cbd6578

2 files changed

Lines changed: 204 additions & 2 deletions

File tree

‎Framework/deploy_handler/long_poll_handler.py‎

Lines changed: 203 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,13 @@
44
import traceback
55
import time
66
import random
7+
import json
78
import requests
89
from colorama import Fore
10+
from pathlib import Path
11+
from urllib.parse import urlparse
912

10-
from Framework.Utilities import RequestFormatter
13+
from Framework.Utilities import RequestFormatter, ConfigModule, CommonUtil
1114
from Framework.node_server_state import STATE
1215
from concurrent.futures import ThreadPoolExecutor
1316

@@ -20,6 +23,8 @@ class DeployHandler:
2023

2124
COMMAND_DONE = b"DONE"
2225
COMMAND_CANCEL = b"CANCEL"
26+
COMMAND_KEY_REQUEST = b"KEY_REQUEST"
27+
COMMAND_PRIVATE_KEY = b"PRIVATE_KEY"
2328
ERROR_PREFIX = b"error"
2429

2530
def __init__(
@@ -50,6 +55,16 @@ def on_message(self, message) -> bool:
5055
# Run cancelled by the user/service.
5156
self.cancel_callback()
5257
return False
58+
59+
elif message.startswith(b'{"command":"KEY_REQUEST"'):
60+
# Handle key request from server
61+
self.handle_key_request(message)
62+
return False
63+
64+
elif message.startswith(b'{"command":"PRIVATE_KEY"'):
65+
# Handle private key received from server
66+
self.handle_private_key(message)
67+
return False
5368

5469
self.response_callback(message)
5570
return False
@@ -60,6 +75,193 @@ def on_error(self, error) -> None:
6075
if self.backoff_time < 6:
6176
self.backoff_time += 1
6277

78+
def handle_key_request(self, message: bytes) -> None:
79+
"""
80+
Handle KEY_REQUEST command.
81+
Server is asking this node to share its private key that matches the provided public key.
82+
"""
83+
try:
84+
payload = json.loads(message)
85+
request_id = payload["request_id"]
86+
public_key_pem = payload["public_key"]
87+
receiver_node_ids = payload["receiver_node_ids"]
88+
89+
print(f"[key-request] Received request {request_id} to share key with {receiver_node_ids}")
90+
91+
# Find the private key that matches the provided public key
92+
private_key_pem = self.get_private_key_matching_public_key(public_key_pem)
93+
94+
if private_key_pem:
95+
# Respond to the key request immediately - no need for separate distribute call!
96+
self.respond_to_key_request(request_id, private_key_pem)
97+
else:
98+
print("[key-request] ERROR: No private key found matching the provided public key")
99+
except Exception as e:
100+
print(f"[key-request] Error handling key request: {e}")
101+
traceback.print_exc()
102+
103+
def handle_private_key(self, message: bytes) -> None:
104+
"""
105+
Handle PRIVATE_KEY command.
106+
Server is sending us a private key.
107+
"""
108+
try:
109+
payload = json.loads(message)
110+
key_id = payload["key_id"]
111+
private_key_pem = payload["private_key"]
112+
113+
print(f"[private-key] Received key {key_id}")
114+
115+
# Save the private key securely
116+
self.save_private_key(key_id, private_key_pem)
117+
118+
print(f"[private-key] Key {key_id} saved successfully")
119+
except Exception as e:
120+
print(f"[private-key] Error handling private key: {e}")
121+
traceback.print_exc()
122+
123+
def get_private_key_matching_public_key(self, public_key_pem: str) -> str | None:
124+
"""
125+
Find and retrieve the private key that corresponds to the provided public key.
126+
Returns the private key in PEM format, or None if no matching key exists.
127+
"""
128+
try:
129+
from settings import ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR
130+
from cryptography.hazmat.primitives import serialization
131+
132+
key_folder = Path(ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR)
133+
if not key_folder.exists():
134+
print("[key-request] Key folder does not exist")
135+
return None
136+
137+
# Parse the provided public key
138+
try:
139+
target_public_key = serialization.load_pem_public_key(
140+
public_key_pem.encode('utf-8')
141+
)
142+
target_public_key_bytes = target_public_key.public_bytes(
143+
encoding=serialization.Encoding.PEM,
144+
format=serialization.PublicFormat.SubjectPublicKeyInfo
145+
)
146+
except Exception as e:
147+
print(f"[key-request] Error parsing provided public key: {e}")
148+
return None
149+
150+
# Get all PEM files
151+
pem_files = list(key_folder.glob("*.pem"))
152+
153+
if not pem_files:
154+
print("[key-request] No private key files found")
155+
return None
156+
157+
# Search through all keys to find matching private key
158+
for pem_file in pem_files:
159+
try:
160+
with open(pem_file, 'rb') as f:
161+
private_key_bytes = f.read()
162+
163+
# Load the private key
164+
private_key = serialization.load_pem_private_key(
165+
private_key_bytes,
166+
password=None
167+
)
168+
169+
# Extract its public key
170+
derived_public_key = private_key.public_key()
171+
derived_public_key_bytes = derived_public_key.public_bytes(
172+
encoding=serialization.Encoding.PEM,
173+
format=serialization.PublicFormat.SubjectPublicKeyInfo
174+
)
175+
176+
# Compare public keys
177+
if derived_public_key_bytes == target_public_key_bytes:
178+
print(f"[key-request] Found matching private key: {pem_file.name}")
179+
return private_key_bytes.decode('utf-8')
180+
181+
except Exception as e:
182+
print(f"[key-request] Error reading key file {pem_file.name}: {e}")
183+
continue
184+
185+
print("[key-request] No private key matches the provided public key")
186+
return None
187+
188+
except Exception as e:
189+
print(f"[key-request] Error searching for matching private key: {e}")
190+
traceback.print_exc()
191+
return None
192+
193+
def save_private_key(self, key_id: str, private_key_pem: str) -> None:
194+
"""
195+
Save the received private key to encrypted storage.
196+
This key is saved to the same location where generated keys are stored.
197+
"""
198+
try:
199+
from settings import ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR
200+
from cryptography.hazmat.primitives import serialization
201+
from datetime import datetime as dt
202+
203+
key_folder = Path(ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR)
204+
key_folder.mkdir(parents=True, exist_ok=True)
205+
206+
# Validate the private key
207+
private_key = serialization.load_pem_private_key(
208+
private_key_pem.encode('utf-8'),
209+
password=None,
210+
)
211+
212+
# Save with descriptive filename
213+
timestamp = dt.now().strftime("%Y%m%d_%H%M%S")
214+
key_filename = f"received_key_{key_id}_{timestamp}.pem"
215+
key_path = key_folder / key_filename
216+
217+
# Save the key
218+
with open(key_path, 'wb') as f:
219+
f.write(private_key.private_bytes(
220+
encoding=serialization.Encoding.PEM,
221+
format=serialization.PrivateFormat.PKCS8,
222+
encryption_algorithm=serialization.NoEncryption()
223+
))
224+
225+
print(f"[private-key] Saved to: {key_path}")
226+
except Exception as e:
227+
print(f"[private-key] Error saving private key: {e}")
228+
traceback.print_exc()
229+
230+
def respond_to_key_request(self, request_id: str, private_key_pem: str) -> None:
231+
"""
232+
Respond to a key request - server will automatically distribute to receivers.
233+
"""
234+
try:
235+
server_url = urlparse(
236+
ConfigModule.get_config_value("Authentication", "server_address")
237+
)
238+
api_url = f"{server_url.scheme}://{server_url.netloc}/zsvc/deploy/v1/respond-to-key-request"
239+
240+
# Get node ID
241+
node_id = CommonUtil.MachineInfo().getLocalUser().lower()
242+
243+
response = RequestFormatter.request(
244+
"post",
245+
api_url,
246+
json_data={
247+
"request_id": request_id,
248+
"donor_node_id": node_id,
249+
"private_key": private_key_pem
250+
},
251+
verify=False
252+
)
253+
254+
if response.ok:
255+
result = response.json()
256+
success_count = result.get('success_count', 0)
257+
print(f"[key-request] Successfully distributed key to {success_count} receiver nodes")
258+
else:
259+
print(f"[key-request] Failed to respond to key request: {response.status_code}")
260+
print(f"[key-request] Response: {response.text}")
261+
except Exception as e:
262+
print(f"[key-request] Error responding to key request: {e}")
263+
traceback.print_exc()
264+
63265
def run(self, host: str) -> None:
64266
reconnect = False
65267
server_online = False

‎settings.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
# BASE_DIR or PROJECT_ROOT or Zeuz_Python_Node dir
88
PROJECT_ROOT = Path(__file__).parent
99

10-
ZEUZ_NODE_ARTIFACTS_DIR = Path.home() / "zeuz"
10+
ZEUZ_NODE_ARTIFACTS_DIR = Path.home() / ".zeuz"
1111
ZEUZ_NODE_DOWNLOADS_DIR = ZEUZ_NODE_ARTIFACTS_DIR / "zeuz_node_downloads"
1212
ZEUZ_NODE_PRIVATE_RSA_KEYS_DIR = ZEUZ_NODE_ARTIFACTS_DIR / "rsa_private_keys"
1313

0 commit comments

Comments
 (0)