Merge branch 'main' of https://git.mamahaha.work/sangge/mimajingsai
This commit is contained in:
commit
4eabfb91de
@ -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")
|
||||
|
||||
@ -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()
|
||||
|
||||
37
src/node.py
37
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)
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user