diff --git a/src/client.py b/src/client.py index a40962d..c9d270a 100644 --- a/src/client.py +++ b/src/client.py @@ -9,6 +9,7 @@ from pydantic import BaseModel import socket import random import time +import base64 @asynccontextmanager @@ -28,7 +29,6 @@ def init(): # load config from config file init_config() - get_node_list(2, server_address) # type: ignore @@ -84,6 +84,9 @@ def init_config(): # execute on exit def clean_env(): + with sqlite3.connect("client.db") as db: + db.execute("DELETE FROM node") + db.commit() print("Exit app") @@ -97,6 +100,7 @@ class C(BaseModel): Tuple: Tuple[capsule, int] ip: str + # receive messages from nodes @app.post("/receive_messages") async def receive_messages(message: C): @@ -108,8 +112,8 @@ async def receive_messages(message: C): return: status_code """ - C_tuple = message.Tuple - + a,b = message.Tuple + C_tuple = (a, b.to_bytes(32)) ip = message.ip if not C_tuple or not ip: raise HTTPException(status_code=400, detail="Invalid input data") @@ -145,7 +149,7 @@ async def receive_messages(message: C): async def check_merge(ct: int, ip: str): global sk, pk, node_response, message 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. cursor = db.execute( """ SELECT capsule, ct @@ -167,8 +171,8 @@ async def check_merge(ct: int, ip: str): (ip), ) result = cursor.fetchall() - pk_sender, T = result[0] # result[0] = (pk, threshold) - + pk_sender, T = result[0] # result[0] = (pk, threshold) + if len(cfrag_cts) >= T: cfrags = mergecfrag(cfrag_cts) message = DecryptFrags(sk, pk, pk_sender, cfrags) # type: ignore @@ -181,29 +185,32 @@ async def send_messages( ): global pk, sk id_list = [] + # calculate id of nodes for node_ip in node_ips: + node_ip = node_ip[0] ip_parts = node_ip.split(".") id = 0 for i in range(4): id += int(ip_parts[i]) << (24 - (8 * i)) id_list.append(id) - # generate rk rk_list = GenerateReKey(sk, pk_B, len(node_ips), shreshold, tuple(id_list)) # type: ignore - - capsule_ct = Encrypt(pk, message) # type: ignore + + capsule, ct = Encrypt(pk, message) # type: ignore + capsule_ct = (capsule, int.from_bytes(ct)) for i in range(len(node_ips)): - url = "http://" + node_ips[i] + ":8001" + "/user_src?message" - + url = "http://" + node_ips[i][0] + ":8001" + "/user_src" payload = { "source_ip": local_ip, "dest_ip": dest_ip, "capsule_ct": capsule_ct, "rk": rk_list[i], } + print(payload) response = requests.post(url, json=payload) + if response.status_code == 200: print(f"send to {node_ips[i]} successful") return 0 @@ -229,8 +236,8 @@ async def request_message(i_m: Request_Message): # dest_ip = dest_ip.split(":")[0] message_name = i_m.message_name source_ip = get_own_ip() - dest_port = "8003" - url = "http://" + dest_ip + ":" + dest_port + "/recieve_request?i_m" + dest_port = "8002" + url = "http://" + dest_ip + ":" + dest_port + "/recieve_request" payload = { "dest_ip": dest_ip, "message_name": message_name, @@ -238,17 +245,33 @@ async def request_message(i_m: Request_Message): "pk": pk, } try: - response = requests.post(url, json=payload) + response = requests.post(url, json=payload, timeout=3) + print(response.text) except: print("can't post") return {"message": "can't post"} + if response.status_code == 200: + data = response.json() + public_key = int(data["public_key"]) + threshold = int(data["threshold"]) + with sqlite3.connect("client.db") as db: + db.execute( + """ + INSERT INTO senderinfo + (public_key, threshold) + VALUES + (?, ?) + """, + (public_key, threshold), + ) try: if response.status_code == 200: data = response.json() public_key = int(data["public_key"]) threshold = int(data["threshold"]) + print(data) with sqlite3.connect("client.db") as db: db.execute( """ @@ -264,10 +287,10 @@ async def request_message(i_m: Request_Message): return {"message": "Database Error"} # wait 10s to recieve message from nodes - for _ in range(10): + for _ in range(3): if node_response: data = message - + # reset message and node_response message = b"" node_response = False @@ -286,7 +309,8 @@ async def recieve_request(i_m: IP_Message): if source_ip != i_m.dest_ip: return HTTPException(status_code=400, detail="Wrong ip") dest_ip = i_m.source_ip - threshold = random.randrange(1, 2) + # threshold = random.randrange(1, 2) + threshold = 2 own_public_key = pk pk_B = i_m.pk @@ -300,18 +324,18 @@ async def recieve_request(i_m: IP_Message): (threshold,), ) node_ips = cursor.fetchall() - + # message name message = b"hello world" + random.randbytes(8) - + # send message to nodes - await send_messages(tuple(node_ips), message, dest_ip, pk_B, threshold) + 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: - ip = os.environ.get("HOST_IP", "IP not set") return ip @@ -319,7 +343,7 @@ def get_own_ip() -> str: # get node list from central server def get_node_list(count: int, server_addr: str): url = "http://" + server_addr + "/server/send_nodes_list?count=" + str(count) - response = requests.get(url,timeout=3) + response = requests.get(url, timeout=3) # Checking the response if response.status_code == 200: print("Success get node list") diff --git a/src/client_cli.py b/src/client_cli.py index 19fef3d..1846f52 100644 --- a/src/client_cli.py +++ b/src/client_cli.py @@ -1,15 +1,14 @@ import argparse import requests + def send_post_request(ip_addr, message_name): - url = f"http://localhost:20234/request_message/?i_m" - data = { - "dest_ip": ip_addr, - "message_name": message_name - } + url = f"http://localhost:8002/request_message" + data = {"dest_ip": ip_addr, "message_name": message_name} response = requests.post(url, json=data) return response.text + def main(): parser = argparse.ArgumentParser(description="Send POST request to a specified IP.") parser.add_argument("ip_addr", help="IP address to send request to.") @@ -19,5 +18,6 @@ def main(): response = send_post_request(args.ip_addr, args.message_name) print(response) + if __name__ == "__main__": main() diff --git a/src/node.py b/src/node.py index d75b1f9..0644121 100644 --- a/src/node.py +++ b/src/node.py @@ -6,6 +6,8 @@ import asyncio from pydantic import BaseModel from tpre import * import os +from typing import Any, Tuple +import base64 @asynccontextmanager @@ -33,7 +35,7 @@ def send_ip(): url = server_address + "/get_node?ip=" + ip # ip = get_local_ip() # type: ignore global id - id = requests.get(url,timeout=3) + id = requests.get(url, timeout=3) # 用环境变量获取本机ip @@ -63,16 +65,23 @@ async def send_heartbeat_internal() -> None: while True: # print('successful send my_heart') try: - folderol = requests.get(url,timeout=3) + folderol = requests.get(url, timeout=3) except: print("Central server error") - + # 删除超时的节点（假设你有一个异步的数据库操作函数） await asyncio.sleep(timeout) +class Req(BaseModel): + source_ip: str + dest_ip: str + capsule_ct: Tuple[capsule, int] + rk: Any + + @app.post("/user_src") # 接收用户1发送的信息 -async def receive_user_src_message(message: Request): +async def user_src(message: Req): global client_ip_src, client_ip_des # kfrag , capsule_ct ,client_ip_src , client_ip_des = json_data[] # 看梁俊勇 """ @@ -83,14 +92,14 @@ async def receive_user_src_message(message: Request): "rk": rk_list[i], } """ + source_ip = message.source_ip + dest_ip = message.dest_ip + capsule, ct = message.capsule_ct + capsule_ct = (capsule, ct.to_bytes(32)) + rk = message.rk - data = await message.json() - source_ip = data.get("source_ip") - dest_ip = data.get("dest_ip") - capsule_ct = data.get("capsule_ct") - rk = data.get("rk") - - processed_message = ReEncrypt(rk, capsule_ct) + a, b = ReEncrypt(rk, capsule_ct) + processed_message = (a, int.from_bytes(b)) await send_user_des_message(source_ip, dest_ip, processed_message) return HTTPException(status_code=200, detail="message recieved") @@ -100,12 +109,12 @@ async def send_user_des_message(source_ip: str, dest_ip: str, re_message): # # 发送 HTTP POST 请求 response = requests.post( - "http://" + dest_ip + "/receive_messages?message", json=data + "http://" + dest_ip + ":8002" + "/receive_messages", json=data ) - print(response) + print(response.text) if __name__ == "__main__": import uvicorn # pylint: disable=e0401 - uvicorn.run("node:app", host="0.0.0.0", port=8001, reload=False) + uvicorn.run("node:app", host="0.0.0.0", port=8001, reload=True) diff --git a/src/server.py b/src/server.py index 46b0fbf..6060f17 100644 --- a/src/server.py +++ b/src/server.py @@ -6,34 +6,41 @@ import sqlite3 import asyncio import time + @asynccontextmanager async def lifespan(app: FastAPI): init() yield clean_env() -app = FastAPI(lifespan = lifespan) + +app = FastAPI(lifespan=lifespan) # 连接到数据库（如果数据库不存在，则会自动创建） -conn = sqlite3.connect('server.db') +conn = sqlite3.connect("server.db") # 创建游标对象，用于执行SQL语句 cursor = conn.cursor() # 创建表: id: int; ip: TEXT -cursor.execute('''CREATE TABLE IF NOT EXISTS nodes ( +cursor.execute( + """CREATE TABLE IF NOT EXISTS nodes ( id INTEGER PRIMARY KEY AUTOINCREMENT, ip TEXT NOT NULL, last_heartbeat INTEGER - )''') + )""" +) + def init(): asyncio.create_task(receive_heartbeat_internal()) + def clean_env(): clear_database() # 关闭游标和连接 cursor.close() conn.close() + @app.get("/server/show_nodes") async def show_nodes() -> list: nodes_list = [] @@ -44,37 +51,42 @@ async def show_nodes() -> list: nodes_list.append(row) return nodes_list + @app.get("/server/get_node") async def get_node(ip: str) -> int: - ''' - 中心服务器与节点交互, 节点发送ip, 中心服务器接收ip存入数据库并将ip转换为int作为节点id返回给节点 - params: - ip: node ip - return: - id: ip按点分割成四部分, 每部分转二进制后拼接再转十进制作为节点id - ''' + """ + 中心服务器与节点交互, 节点发送ip, 中心服务器接收ip存入数据库并将ip转换为int作为节点id返回给节点 + params: + ip: node ip + return: + id: ip按点分割成四部分, 每部分转二进制后拼接再转十进制作为节点id + """ ip_parts = ip.split(".") ip_int = 0 for i in range(4): ip_int += int(ip_parts[i]) << (24 - (8 * i)) - + # 获取当前时间 current_time = int(time.time()) # 插入数据 - cursor.execute("INSERT INTO nodes (id, ip, last_heartbeat) VALUES (?, ?, ?)", (ip_int, ip, current_time)) + cursor.execute( + "INSERT INTO nodes (id, ip, last_heartbeat) VALUES (?, ?, ?)", + (ip_int, ip, current_time), + ) conn.commit() return ip_int + @app.get("/server/delete_node") async def delete_node(ip: str) -> None: - ''' + """ param: ip: 待删除节点的ip地址 return: None - ''' + """ # 查询要删除的节点 cursor.execute("SELECT * FROM nodes WHERE ip=?", (ip,)) row = cursor.fetchone() @@ -86,12 +98,16 @@ async def delete_node(ip: str) -> None: else: print(f"Node with IP {ip} not found.") + # 接收节点心跳包 @app.get("/server/heartbeat") async def receive_heartbeat(ip: str): - cursor.execute("UPDATE nodes SET last_heartbeat = ? WHERE ip = ?", (time.time(), ip)) - return {"status": "received"} - + cursor.execute( + "UPDATE nodes SET last_heartbeat = ? WHERE ip = ?", (time.time(), ip) + ) + return {"status": "received"} + + async def receive_heartbeat_internal(): while 1: timeout = 7 @@ -100,32 +116,34 @@ async def receive_heartbeat_internal(): conn.commit() await asyncio.sleep(timeout) + @app.get("/server/send_nodes_list") async def send_nodes_list(count: int) -> list: - ''' + """ 中心服务器与客户端交互, 客户端发送所需节点个数, 中心服务器从数据库中顺序取出节点封装成list格式返回给客户端 - params: - count: 所需节点个数 - return: + params: + count: 所需节点个数 + return: nodes_list: list - ''' + """ nodes_list = [] # 查询数据库中的节点数据 cursor.execute("SELECT * FROM nodes LIMIT ?", (count,)) rows = cursor.fetchall() - for row in rows: id, ip, last_heartbeat = row nodes_list.append(ip) return nodes_list + # @app.get("/server/clear_database") def clear_database() -> None: cursor.execute("DELETE FROM nodes") conn.commit() + if __name__ == "__main__": import uvicorn # pylint: disable=e0401