Merge pull request 'main' (#17) from sangge/mimajingsai:main into main

Reviewed-on: dqy/mimajingsai#17

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")

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")