diff --git a/src/client.py b/src/client.py index d557dca..d80688d 100644 --- a/src/client.py +++ b/src/client.py @@ -28,7 +28,6 @@ def init(): # load config from config file init_config() - get_node_list(2, server_address) # type: ignore @@ -100,6 +99,7 @@ class C(BaseModel): Tuple: Tuple[capsule, int] ip: str + # receive messages from nodes @app.post("/receive_messages") async def receive_messages(message: C): @@ -148,7 +148,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 @@ -170,8 +170,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 @@ -184,6 +184,7 @@ async def send_messages( ): global pk, sk id_list = [] + print(node_ips) # calculate id of nodes for node_ip in node_ips: ip_parts = node_ip.split(".") @@ -191,10 +192,10 @@ async def send_messages( 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 for i in range(len(node_ips)): @@ -246,6 +247,20 @@ async def request_message(i_m: Request_Message): 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: @@ -271,7 +286,7 @@ async def request_message(i_m: Request_Message): for _ in range(3): if node_response: data = message - + # reset message and node_response message = b"" node_response = False @@ -290,7 +305,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 @@ -304,18 +320,17 @@ 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} return response def get_own_ip() -> str: - ip = os.environ.get("HOST_IP", "IP not set") return ip @@ -323,7 +338,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/node.py b/src/node.py index d75b1f9..99ff06c 100644 --- a/src/node.py +++ b/src/node.py @@ -33,7 +33,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,10 +63,10 @@ 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) 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