sangge/tpre-python
2
1
Fork 1
Code Issues 6 Pull Requests Packages Projects 1 Activity

feat: add debug info

Browse Source
This commit is contained in:
Smart-SangGe 2023-10-27 00:08:01 +08:00
parent bfb5c27bcb
commit 9599ef4321
1 changed files with 126 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

178
src/client.py
View File

@ -80,11 +80,21 @@ def init_db():
def init_config(): def init_config():
import configparser import configparser
print("Starting function: init_config")
global server_address global server_address
config = configparser.ConfigParser() config = configparser.ConfigParser()
print("Attempting to read client.ini...")
config.read("client.ini") config.read("client.ini")
server_address = config["settings"]["server_address"] if "settings" in config and "server_address" in config["settings"]:
server_address = config["settings"]["server_address"]
print(f"Config loaded successfully. Server address: {server_address}")
else:
print("Error: 'settings' section or 'server_address' key not found in client.ini")
print("Function init_config executed successfully!")
# execute on exit # execute on exit
@ -122,18 +132,25 @@ async def receive_messages(message: C):
status_code status_code
""" """
print("Starting function: receive_messages")
if not message.Tuple or not message.ip: if not message.Tuple or not message.ip:
print("Invalid input data")
raise HTTPException(status_code=400, detail="Invalid input data") raise HTTPException(status_code=400, detail="Invalid input data")
C_capsule, C_ct = message.Tuple C_capsule, C_ct = message.Tuple
ip = message.ip ip = message.ip
print(f"Received message: Capsule = {C_capsule}, C_ct = {C_ct}, IP = {ip}")
# Serialization # Serialization
print("Serializing the capsule...")
bin_C_capsule = pickle.dumps(C_capsule) bin_C_capsule = pickle.dumps(C_capsule)
print("Serialization successful")
# insert record into database # insert record into database
with sqlite3.connect("client.db") as db: with sqlite3.connect("client.db") as db:
try: try:
print("Attempting to insert data into 'message' table...")
db.execute( db.execute(
""" """
INSERT INTO message INSERT INTO message
@ -144,28 +161,30 @@ async def receive_messages(message: C):
(bin_C_capsule, str(C_ct), ip), (bin_C_capsule, str(C_ct), ip),
) )
db.commit() db.commit()
print("Data insertion successful")
check_merge(C_ct, ip) check_merge(C_ct, ip)
print("check_merge executed successfully")
return HTTPException(status_code=200, detail="Message received") return HTTPException(status_code=200, detail="Message received")
except Exception as e: except Exception as e:
print(f"Error occurred: {e}") print(f"Error occurred: {e}")
db.rollback() db.rollback()
return HTTPException(status_code=400, detail="Database error") return HTTPException(status_code=400, detail="Database error")
print("Function receive_messages executed successfully!")
# check record count # check record count
def check_merge(ct: int, ip: str): def check_merge(ct: int, ip: str):
print("Starting function: check_merge")
global sk, pk, node_response, message global sk, pk, node_response, message
"""
CREATE TABLE IF NOT EXISTS senderinfo (
id INTEGER PRIMARY KEY,
ip TEXT,
pkx TEXT,
pky TEXT,
threshold INTEGER
)
"""
with sqlite3.connect("client.db") as db: with sqlite3.connect("client.db") as db:
# Check if the combination of ct_column and ip_column appears more than once. # Check if the combination of ct_column and ip_column appears more than once.
print("Fetching data from 'message' table...")
cursor = db.execute( cursor = db.execute(
""" """
SELECT capsule, ct SELECT capsule, ct
@ -174,10 +193,11 @@ def check_merge(ct: int, ip: str):
""", """,
(str(ct), ip), (str(ct), ip),
) )
# [(capsule, ct), ...]
cfrag_cts = cursor.fetchall() cfrag_cts = cursor.fetchall()
print(f"Number of records fetched from 'message' table: {len(cfrag_cts)}")
# get _sender_pk # get _sender_pk
print("Fetching sender's public key...")
cursor = db.execute( cursor = db.execute(
""" """
SELECT pkx, pky SELECT pkx, pky
@ -190,8 +210,10 @@ def check_merge(ct: int, ip: str):
try: try:
pkx, pky = result[0] # result[0] = (pkx, pky) pkx, pky = result[0] # result[0] = (pkx, pky)
pk_sender = (int(pkx), int(pky)) pk_sender = (int(pkx), int(pky))
print(f"Successfully fetched sender's public key: {pk_sender}")
except: except:
pk_sender, T = 0, -1 pk_sender, T = 0, -1
print("Failed to fetch sender's public key")
T = 2 T = 2
if len(cfrag_cts) >= T: if len(cfrag_cts) >= T:
@ -200,29 +222,38 @@ def check_merge(ct: int, ip: str):
for i in cfrag_cts: for i in cfrag_cts:
capsule = pickle.loads(i[0]) capsule = pickle.loads(i[0])
temp_cfrag_cts.append((capsule, int(i[1]).to_bytes(32))) temp_cfrag_cts.append((capsule, int(i[1]).to_bytes(32)))
print("Deserialization completed")
cfrags = mergecfrag(temp_cfrag_cts) cfrags = mergecfrag(temp_cfrag_cts)
print("Attempting decryption...")
print("sk:", type(sk)) print("sk:", type(sk))
print("pk:", type(pk)) print("pk:", type(pk))
print("pk_sender:", type(pk_sender)) print("pk_sender:", type(pk_sender))
print("cfrags:", type(cfrags)) print("cfrags:", type(cfrags))
message = DecryptFrags(sk, pk, pk_sender, cfrags) message = DecryptFrags(sk, pk, pk_sender, cfrags)
print("merge success", message) print(f"Decryption successful, message: {message}")
node_response = True node_response = True
print(f"Node response set to: {node_response}")
print("merge:", node_response) else:
print("Insufficient number of cfrag_cts, skipping decryption")
print("Function check_merge executed successfully!")
# send message to node # send message to node
async def send_messages( async def send_messages(
node_ips: tuple[str, ...], message: bytes, dest_ip: str, pk_B: point, shreshold: int node_ips: tuple[str, ...], message: bytes, dest_ip: str, pk_B: point, shreshold: int
): ):
print("Starting function: send_messages")
global pk, sk global pk, sk
id_list = [] id_list = []
# calculate id of nodes # calculate id of nodes
print("Calculating ID of nodes...")
for node_ip in node_ips: for node_ip in node_ips:
node_ip = node_ip[0] node_ip = node_ip[0]
ip_parts = node_ip.split(".") ip_parts = node_ip.split(".")
@ -230,8 +261,12 @@ async def send_messages(
for i in range(4): for i in range(4):
id += int(ip_parts[i]) << (24 - (8 * i)) id += int(ip_parts[i]) << (24 - (8 * i))
id_list.append(id) id_list.append(id)
print(f"Calculated IDs: {id_list}")
# generate rk # generate rk
print("Generating rekey...")
rk_list = GenerateReKey(sk, pk_B, len(node_ips), shreshold, tuple(id_list)) # type: ignore rk_list = GenerateReKey(sk, pk_B, len(node_ips), shreshold, tuple(id_list)) # type: ignore
print(f"Generated ReKey: {rk_list}")
capsule, ct = Encrypt(pk, message) # type: ignore capsule, ct = Encrypt(pk, message) # type: ignore
# capsule_ct = (capsule, int.from_bytes(ct)) # capsule_ct = (capsule, int.from_bytes(ct))
@ -245,14 +280,20 @@ async def send_messages(
"ct": int.from_bytes(ct), "ct": int.from_bytes(ct),
"rk": rk_list[i], "rk": rk_list[i],
} }
print(f"Sending payload to {url}:")
print(json.dumps(payload)) print(json.dumps(payload))
response = requests.post(url, json=payload) response = requests.post(url, json=payload)
if response.status_code == 200: if response.status_code == 200:
print(f"send to {node_ips[i]} successful") print(f"send to {node_ips[i]} successful")
else:
print(f"send to {node_ips[i]} failed with status code {response.status_code}")
print("Function send_messages executed successfully!")
return 0 return 0
class IP_Message(BaseModel): class IP_Message(BaseModel):
dest_ip: str dest_ip: str
message_name: str message_name: str
@ -268,6 +309,8 @@ class Request_Message(BaseModel):
# request message from others # request message from others
@app.post("/request_message") @app.post("/request_message")
async def request_message(i_m: Request_Message): async def request_message(i_m: Request_Message):
print("Starting function: request_message")
global message, node_response, pk global message, node_response, pk
dest_ip = i_m.dest_ip dest_ip = i_m.dest_ip
# dest_ip = dest_ip.split(":")[0] # dest_ip = dest_ip.split(":")[0]
@ -281,64 +324,88 @@ async def request_message(i_m: Request_Message):
"source_ip": source_ip, "source_ip": source_ip,
"pk": pk, "pk": pk,
} }
print(f"Requesting message from: {url}")
try: try:
response = requests.post(url, json=payload, timeout=1) response = requests.post(url, json=payload, timeout=1)
# print("menxian and pk", response.text) print(f"Response received from {url}: {response.text}")
except requests.Timeout: except requests.Timeout as e:
print(f"Request to {url} timed out!")
print("can't post") print("can't post")
# content = {"message": "post timeout", "error": str(e)}
# return JSONResponse(content, status_code=400)
# wait 3s to receive message from nodes # wait 2s to receive message from nodes
for _ in range(10): for _ in range(10):
print("wait:", node_response) print(f"Waiting for response... (iteration {_ + 1})")
print("Current node_response:", node_response)
if node_response: if node_response:
data = message data = message
# reset message and node_response # reset message and node_response
print("Resetting message and node_response...")
message = b"" message = b""
node_response = False node_response = False
# return message to frontend # return message to frontend
print("Returning message to frontend:", str(data))
return {"message": str(data)} return {"message": str(data)}
await asyncio.sleep(0.2) await asyncio.sleep(0.2)
print("Timeout occurred while waiting for response.")
content = {"message": "receive timeout"} content = {"message": "receive timeout"}
return JSONResponse(content, status_code=400) return JSONResponse(content, status_code=400)
# receive request from others
@app.post("/receive_request") # request message from others
async def receive_request(i_m: IP_Message): @app.post("/request_message")
global pk async def request_message(i_m: Request_Message):
print("Starting function: request_message")
global message, node_response, pk
dest_ip = i_m.dest_ip
# dest_ip = dest_ip.split(":")[0]
message_name = i_m.message_name
source_ip = get_own_ip() source_ip = get_own_ip()
if source_ip != i_m.dest_ip: dest_port = "8002"
return HTTPException(status_code=400, detail="Wrong ip") url = "http://" + dest_ip + ":" + dest_port + "/receive_request"
dest_ip = i_m.source_ip payload = {
# threshold = random.randrange(1, 2) "dest_ip": dest_ip,
threshold = 2 "message_name": message_name,
own_public_key = pk "source_ip": source_ip,
pk_B = i_m.pk "pk": pk,
}
with sqlite3.connect("client.db") as db: print(f"Requesting message from: {url}")
cursor = db.execute( try:
""" response = requests.post(url, json=payload, timeout=1)
SELECT nodeip print(f"Response received from {url}: {response.text}")
FROM node
LIMIT ?
""",
(threshold,),
)
node_ips = cursor.fetchall()
# message name except requests.Timeout as e:
message = b"hello world" + random.randbytes(8) print(f"Request to {url} timed out!")
print("can't post")
# wait 2s to receive message from nodes
for _ in range(10):
print(f"Waiting for response... (iteration {_ + 1})")
print("Current node_response:", node_response)
if node_response:
data = message
# reset message and node_response
print("Resetting message and node_response...")
message = b""
node_response = False
# return message to frontend
print("Returning message to frontend:", str(data))
return {"message": str(data)}
await asyncio.sleep(0.2)
print("Timeout occurred while waiting for response.")
content = {"message": "receive timeout"}
return JSONResponse(content, status_code=400)
# send message to nodes
await send_messages(tuple(node_ips), message, dest_ip, pk_B, threshold)
response = {"threshold": threshold, "public_key": own_public_key}
print("###############RESPONSE = ", response)
return response
def get_own_ip() -> str: def get_own_ip() -> str:
@ -385,16 +452,21 @@ class pk_model(BaseModel):
pky: str pky: str
ip: str ip: str
# recieve pk from frontend
@app.post("/recieve_pk") @app.post("/recieve_pk")
async def recieve_pk(pk: pk_model): async def recieve_pk(pk: pk_model):
print("Starting function: recieve_pk")
pkx = pk.pkx pkx = pk.pkx
pky = pk.pky pky = pk.pky
dest_ip = pk.ip dest_ip = pk.ip
print(f"Received pkx: {pkx}, pky: {pky}, IP: {dest_ip}")
try: try:
threshold = 2 threshold = 2
print("Connecting to client.db...")
with sqlite3.connect("client.db") as db: with sqlite3.connect("client.db") as db:
print("Connected to client.db, inserting data...")
db.execute( db.execute(
""" """
INSERT INTO senderinfo INSERT INTO senderinfo
@ -404,13 +476,15 @@ async def recieve_pk(pk: pk_model):
""", """,
(str(dest_ip), pkx, pky, threshold), (str(dest_ip), pkx, pky, threshold),
) )
print("Data inserted successfully!")
except Exception as e: except Exception as e:
# raise error # raise error
print("Database error") print("Database error:", str(e))
content = {"message": "Database Error", "error": str(e)} content = {"message": "Database Error", "error": str(e)}
return JSONResponse(content, status_code=400) return JSONResponse(content, status_code=400)
return {"message": "save pk in database"}
print("Function recieve_pk executed successfully!")
return {"message": "save pk in database"}
pk = (0, 0) pk = (0, 0)
sk = 0 sk = 0