diff --git a/src/client.py b/src/client.py index 28e4f69..7044b5f 100644 --- a/src/client.py +++ b/src/client.py @@ -5,6 +5,10 @@ from typing import Tuple from tpre import * import sqlite3 from contextlib import asynccontextmanager +from pydantic import BaseModel +import socket +import random +import time @asynccontextmanager @@ -19,6 +23,8 @@ app = FastAPI(lifespan=lifespan) pk = point sk = int server_address = str +node_response = False +message = bytes def init(): @@ -26,7 +32,7 @@ def init(): init_db() pk, sk = GenerateKeyPair() init_config() - get_node_list(6, server_address) # type: ignore + # get_node_list(6, server_address) # type: ignore def init_db(): @@ -68,6 +74,7 @@ def init_db(): print("Init Database Successful") +# load config from config file def init_config(): import configparser @@ -89,9 +96,14 @@ async def read_root(): return {"message": "Hello, World!"} +class C(BaseModel): + Tuple: Tuple[capsule, int] + ip: str + + # receive messages from node @app.post("/receive_messages") -async def receive_messages(C: Tuple[capsule, int], ip: str): +async def receive_messages(message: C): """ receive capsule and ip from nodes params: @@ -100,22 +112,32 @@ async def receive_messages(C: Tuple[capsule, int], ip: str): return: status_code """ - if not C or not ip: + C_tuple = message.Tuple + + ip = message.ip + if not C_tuple or not ip: raise HTTPException(status_code=400, detail="Invalid input data") - capsule, ct = C - if not Checkcapsule(capsule): + C_capsule = C_tuple[0] + C_ct = C_tuple[1] + + if not Checkcapsule(C_capsule): raise HTTPException(status_code=400, detail="Invalid capsule") # insert record into database with sqlite3.connect("message.db") as db: try: db.execute( - "INSERT INTO message (capsule_column, ct_column, ip_column) VALUES (?, ?, ?)", - (capsule, ct, ip), + """ + INSERT INTO message + (capsule_column, ct_column, ip_column) + VALUES + (?, ?, ?) + """, + (C_capsule, C_ct, ip), ) db.commit() - await check_merge(db, ct, ip) + await check_merge(db, C_ct, ip) return HTTPException(status_code=200, detail="Message received") except Exception as e: print(f"Error occurred: {e}") @@ -125,7 +147,7 @@ async def receive_messages(C: Tuple[capsule, int], ip: str): # check record count async def check_merge(db, ct: int, ip: str): - global sk, pk + global sk, pk, node_response, message # Check if the combination of ct_column and ip_column appears more than once. cursor = db.execute( """ @@ -151,19 +173,75 @@ async def check_merge(db, ct: int, ip: str): pk_sender, T = result[0] if len(cfrag_cts) >= T: cfrags = mergecfrag(cfrag_cts) - m = DecryptFrags(sk, pk, pk_sender, cfrags) # type: ignore + message = DecryptFrags(sk, pk, pk_sender, cfrags) # type: ignore + node_response = True # send message to node -@app.post("/send_message") -async def send_message(ip: tuple[str, ...]): +def send_message(ip: tuple[str, ...]): return 0 +class IP_Message(BaseModel): + dest_ip: str + message_name: str + source_ip: str + + # request message from others @app.post("/request_message") -async def request_message(ip): - return 0 +async def request_message(i_m: IP_Message): + global message, node_response + dest_ip = i_m.dest_ip + message_name = i_m.message_name + source_ip = get_own_ip() + dest_port = "8003" + url = "http://" + dest_ip + dest_port + "/recieve_request" + payload = {"dest_ip": dest_ip, "message_name": message_name, "source_ip": source_ip} + response = requests.post(url, json=payload) + 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), + ) + + # wait to recieve message from nodes + for _ in range(10): + if node_response: + data = message + message = b"" + # return message to frontend + return {"message": data} + time.sleep(1) + + +# recieve request from others +@app.post("/recieve_request") +async def recieve_request(i_m: IP_Message): + global pk + source_ip = get_own_ip() + 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, 6) + public_key = pk + response = {"threshold": threshold,"public_key": public_key} + return response + + +def get_own_ip() -> str: + hostname = socket.gethostname() + ip = socket.gethostbyname(hostname) + return ip # get node list from central server @@ -194,4 +272,4 @@ def get_node_list(count: int, server_addr: str): if __name__ == "__main__": import uvicorn # pylint: disable=e0401 - uvicorn.run("client:app", host="0.0.0.0", port=8003) + uvicorn.run("client:app", host="0.0.0.0", port=8003, reload="True") diff --git a/src/server.py b/src/server.py index 30481fe..d030e2c 100644 --- a/src/server.py +++ b/src/server.py @@ -28,7 +28,7 @@ cursor.execute('''CREATE TABLE IF NOT EXISTS nodes ( )''') def init(): - nothing = receive_heartbeat_internal() + task = asyncio.create_task(receive_heartbeat_internal()) def clean_env(): # 关闭游标和连接 @@ -101,7 +101,7 @@ async def receive_heartbeat_internal() -> int: cursor.execute("DELETE FROM nodes WHERE last_heartbeat < ?", (time.time() - timeout,)) conn.commit() print('successful delete') - time.sleep(timeout) + await asyncio.sleep(timeout) return 1 @app.get("/server/send_nodes_list")