Merge branch 'main' of https://git.mamahaha.work/sangge/mimajingsai

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,7 +171,7 @@ 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)
@@ -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"
+    dest_port = "8002"
-    url = "http://" + dest_ip + ":" + dest_port + "/recieve_request?i_m"
+    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,7 +287,7 @@ 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
 
@@ -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
 
@@ -307,11 +331,11 @@ async def recieve_request(i_m: IP_Message):
     # 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:
-    
     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")

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"
+    url = f"http://localhost:8002/request_message"
-    data = {
+    data = {"dest_ip": ip_addr, "message_name": message_name}
-        "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()

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,7 +65,7 @@ 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")
 
@@ -71,8 +73,15 @@ async def send_heartbeat_internal() -> None:
         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()
+    a, b = ReEncrypt(rk, capsule_ct)
-    source_ip = data.get("source_ip")
+    processed_message = (a, int.from_bytes(b))
-    dest_ip = data.get("dest_ip")
-    capsule_ct = data.get("capsule_ct")
-    rk = data.get("rk")
-
-    processed_message = ReEncrypt(rk, capsule_ct)
     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)

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,15 +51,16 @@ 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_parts = ip.split(".")
     ip_int = 0
     for i in range(4):
@@ -62,19 +70,23 @@ async def get_node(ip: str) -> int:
     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,11 +98,15 @@ 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))
+    cursor.execute(
-        return {"status": "received"}     
+        "UPDATE nodes SET last_heartbeat = ? WHERE ip = ?", (time.time(), ip)
+    )
+    return {"status": "received"}
+
 
 async def receive_heartbeat_internal():
     while 1:
@@ -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:
     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