463 lines
12 KiB
Python
463 lines
12 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
import requests
|
|
import os
|
|
from typing import Tuple
|
|
from tpre import *
|
|
import sqlite3
|
|
from contextlib import asynccontextmanager
|
|
from pydantic import BaseModel
|
|
import socket
|
|
import random
|
|
import time
|
|
import base64
|
|
import json
|
|
import pickle
|
|
from fastapi.responses import JSONResponse
|
|
import asyncio
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
init()
|
|
yield
|
|
clean_env()
|
|
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
|
|
def init():
|
|
global pk, sk, server_address
|
|
init_db()
|
|
pk, sk = GenerateKeyPair()
|
|
|
|
# load config from config file
|
|
init_config()
|
|
get_node_list(2, server_address) # type: ignore
|
|
|
|
|
|
def init_db():
|
|
with sqlite3.connect("client.db") as db:
|
|
# message table
|
|
db.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS message (
|
|
id INTEGER PRIMARY KEY,
|
|
capsule BLOB,
|
|
ct TEXT,
|
|
senderip TEXT
|
|
);
|
|
"""
|
|
)
|
|
|
|
# node ip table
|
|
db.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS node (
|
|
id INTEGER PRIMARY KEY,
|
|
nodeip TEXT
|
|
);
|
|
"""
|
|
)
|
|
|
|
# sender info table
|
|
db.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS senderinfo (
|
|
id INTEGER PRIMARY KEY,
|
|
ip TEXT,
|
|
pkx TEXT,
|
|
pky TEXT,
|
|
threshold INTEGER
|
|
)
|
|
"""
|
|
)
|
|
db.commit()
|
|
print("Init Database Successful")
|
|
|
|
|
|
# load config from config file
|
|
def init_config():
|
|
import configparser
|
|
|
|
global server_address
|
|
config = configparser.ConfigParser()
|
|
config.read("client.ini")
|
|
|
|
server_address = config["settings"]["server_address"]
|
|
|
|
|
|
# execute on exit
|
|
def clean_env():
|
|
global message, node_response
|
|
message = b""
|
|
node_response = False
|
|
with sqlite3.connect("client.db") as db:
|
|
db.execute("DELETE FROM node")
|
|
db.execute("DELETE FROM message")
|
|
db.execute("DELETE FROM senderinfo")
|
|
db.commit()
|
|
print("Exit app")
|
|
|
|
|
|
# main page
|
|
@app.get("/")
|
|
async def read_root():
|
|
return {"message": "Hello, World!"}
|
|
|
|
|
|
class C(BaseModel):
|
|
Tuple: Tuple[Tuple[Tuple[int, int], Tuple[int, int], int, Tuple[int, int]], int]
|
|
ip: str
|
|
|
|
|
|
# receive messages from nodes
|
|
@app.post("/receive_messages")
|
|
async def receive_messages(message: C):
|
|
"""
|
|
receive capsule and ip from nodes
|
|
params:
|
|
Tuple: capsule and ct
|
|
ip: sender ip
|
|
return:
|
|
status_code
|
|
"""
|
|
print(f"Received message: {message}")
|
|
|
|
if not message.Tuple or not message.ip:
|
|
print("Invalid input data received.")
|
|
raise HTTPException(status_code=400, detail="Invalid input data")
|
|
|
|
C_capsule, C_ct = message.Tuple
|
|
ip = message.ip
|
|
|
|
# Serialization
|
|
bin_C_capsule = pickle.dumps(C_capsule)
|
|
|
|
# insert record into database
|
|
with sqlite3.connect("client.db") as db:
|
|
try:
|
|
db.execute(
|
|
"""
|
|
INSERT INTO message
|
|
(capsule, ct, senderip)
|
|
VALUES
|
|
(?, ?, ?)
|
|
""",
|
|
(bin_C_capsule, str(C_ct), ip),
|
|
)
|
|
db.commit()
|
|
print("Data inserted successfully into database.")
|
|
check_merge(C_ct, ip)
|
|
return HTTPException(status_code=200, detail="Message received")
|
|
except Exception as e:
|
|
print(f"Error occurred: {e}")
|
|
db.rollback()
|
|
return HTTPException(status_code=400, detail="Database error")
|
|
|
|
|
|
# check record count
|
|
def check_merge(ct: int, ip: str):
|
|
global sk, pk, node_response, message
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS senderinfo (
|
|
id INTEGER PRIMARY KEY,
|
|
ip TEXT,
|
|
pkx TEXT,
|
|
pky TEXT,
|
|
threshold INTEGER
|
|
)
|
|
"""
|
|
with sqlite3.connect("client.db") as db:
|
|
# Check if the combination of ct_column and ip_column appears more than once.
|
|
cursor = db.execute(
|
|
"""
|
|
SELECT capsule, ct
|
|
FROM message
|
|
WHERE ct = ? AND senderip = ?
|
|
""",
|
|
(str(ct), ip),
|
|
)
|
|
# [(capsule, ct), ...]
|
|
cfrag_cts = cursor.fetchall()
|
|
|
|
# get _sender_pk
|
|
cursor = db.execute(
|
|
"""
|
|
SELECT pkx, pky
|
|
FROM senderinfo
|
|
WHERE ip = ?
|
|
""",
|
|
(ip,),
|
|
)
|
|
result = cursor.fetchall()
|
|
try:
|
|
pkx, pky = result[0] # result[0] = (pkx, pky)
|
|
pk_sender = (int(pkx), int(pky))
|
|
except:
|
|
pk_sender, T = 0, -1
|
|
|
|
T = 2
|
|
if len(cfrag_cts) >= T:
|
|
# Deserialization
|
|
temp_cfrag_cts = []
|
|
for i in cfrag_cts:
|
|
capsule = pickle.loads(i[0])
|
|
temp_cfrag_cts.append((capsule, int(i[1]).to_bytes(32)))
|
|
|
|
cfrags = mergecfrag(temp_cfrag_cts)
|
|
|
|
print("sk:", type(sk))
|
|
print("pk:", type(pk))
|
|
print("pk_sender:", type(pk_sender))
|
|
print("cfrags:", type(cfrags))
|
|
message = DecryptFrags(sk, pk, pk_sender, cfrags)
|
|
|
|
print("merge success", message)
|
|
node_response = True
|
|
|
|
print("merge:", node_response)
|
|
|
|
|
|
# send message to node
|
|
async def send_messages(
|
|
node_ips: tuple[str, ...], message: bytes, dest_ip: str, pk_B: point, shreshold: int
|
|
):
|
|
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)
|
|
print(f"Calculated IDs: {id_list}")
|
|
# generate rk
|
|
rk_list = GenerateReKey(sk, pk_B, len(node_ips), shreshold, tuple(id_list)) # type: ignore
|
|
print(f"Generated ReKey list: {rk_list}")
|
|
capsule, ct = Encrypt(pk, message) # type: ignore
|
|
# capsule_ct = (capsule, int.from_bytes(ct))
|
|
print(f"Encrypted message to capsule={capsule}, ct={ct}")
|
|
|
|
for i in range(len(node_ips)):
|
|
url = "http://" + node_ips[i][0] + ":8001" + "/user_src"
|
|
payload = {
|
|
"source_ip": local_ip,
|
|
"dest_ip": dest_ip,
|
|
"capsule": capsule,
|
|
"ct": int.from_bytes(ct),
|
|
"rk": rk_list[i],
|
|
}
|
|
print(f"Sending payload to {url}: {json.dumps(payload)}")
|
|
response = requests.post(url, json=payload)
|
|
|
|
if response.status_code == 200:
|
|
print(f"send to {node_ips[i]} successful")
|
|
else:
|
|
print(
|
|
f"Failed to send to {node_ips[i]}. Response code: {response.status_code}, Response text: {response.text}"
|
|
)
|
|
return 0
|
|
|
|
|
|
class IP_Message(BaseModel):
|
|
dest_ip: str
|
|
message_name: str
|
|
source_ip: str
|
|
pk: Tuple[int, int]
|
|
|
|
|
|
class Request_Message(BaseModel):
|
|
dest_ip: str
|
|
message_name: str
|
|
|
|
|
|
# request message from others
|
|
@app.post("/request_message")
|
|
async def request_message(i_m: Request_Message):
|
|
global message, node_response, pk
|
|
print(
|
|
f"Function 'request_message' called with: dest_ip={i_m.dest_ip}, message_name={i_m.message_name}"
|
|
)
|
|
dest_ip = i_m.dest_ip
|
|
# dest_ip = dest_ip.split(":")[0]
|
|
message_name = i_m.message_name
|
|
source_ip = get_own_ip()
|
|
dest_port = "8002"
|
|
url = "http://" + dest_ip + ":" + dest_port + "/receive_request"
|
|
payload = {
|
|
"dest_ip": dest_ip,
|
|
"message_name": message_name,
|
|
"source_ip": source_ip,
|
|
"pk": pk,
|
|
}
|
|
print(f"Sending request to {url} with payload: {payload}")
|
|
try:
|
|
response = requests.post(url, json=payload, timeout=1)
|
|
print(f"Response received from {url}: {response.text}")
|
|
# print("menxian and pk", response.text)
|
|
|
|
except requests.Timeout:
|
|
print("Timeout error: can't post to the destination.")
|
|
# print("can't post")
|
|
# content = {"message": "post timeout", "error": str(e)}
|
|
# return JSONResponse(content, status_code=400)
|
|
|
|
# wait 3s to receive message from nodes
|
|
for _ in range(10):
|
|
print(f"Waiting for node_response... Current value: {node_response}")
|
|
# print("wait:", node_response)
|
|
if node_response:
|
|
data = message
|
|
print(f"Node response received with message: {data}")
|
|
# reset message and node_response
|
|
message = b""
|
|
node_response = False
|
|
|
|
# return message to frontend
|
|
return {"message": str(data)}
|
|
await asyncio.sleep(0.2)
|
|
print("Timeout while waiting for node_response.")
|
|
content = {"message": "receive timeout"}
|
|
return JSONResponse(content, status_code=400)
|
|
|
|
|
|
# receive request from others
|
|
@app.post("/receive_request")
|
|
async def receive_request(i_m: IP_Message):
|
|
global pk
|
|
print(
|
|
f"Function 'receive_request' called with: dest_ip={i_m.dest_ip}, source_ip={i_m.source_ip}, pk={i_m.pk}"
|
|
)
|
|
source_ip = get_own_ip()
|
|
print(f"Own IP: {source_ip}")
|
|
if source_ip != i_m.dest_ip:
|
|
print("Mismatch in destination IP.")
|
|
return HTTPException(status_code=400, detail="Wrong ip")
|
|
dest_ip = i_m.source_ip
|
|
# threshold = random.randrange(1, 2)
|
|
threshold = 2
|
|
own_public_key = pk
|
|
pk_B = i_m.pk
|
|
print(f"Using own public key: {own_public_key} and received public key: {pk_B}")
|
|
|
|
with sqlite3.connect("client.db") as db:
|
|
cursor = db.execute(
|
|
"""
|
|
SELECT nodeip
|
|
FROM node
|
|
LIMIT ?
|
|
""",
|
|
(threshold,),
|
|
)
|
|
node_ips = cursor.fetchall()
|
|
print(f"Selected node IPs from database: {node_ips}")
|
|
|
|
# message name
|
|
# message_name = i_m.message_name
|
|
# message = xxxxx
|
|
message = b"hello world" + random.randbytes(8)
|
|
print(f"Generated message: {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(f"Sending response: {response}")
|
|
return response
|
|
|
|
|
|
def get_own_ip() -> str:
|
|
ip = os.environ.get("HOST_IP")
|
|
if not ip: # 如果环境变量中没有IP
|
|
try:
|
|
# 从网卡获取IP
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
s.connect(("8.8.8.8", 80)) # 通过连接Google DNS获取IP
|
|
ip = s.getsockname()[0]
|
|
s.close()
|
|
except:
|
|
raise ValueError("Unable to get IP")
|
|
return str(ip)
|
|
|
|
|
|
# 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)
|
|
# Checking the response
|
|
if response.status_code == 200:
|
|
print("Success get node list")
|
|
node_ip = response.text
|
|
node_ip = eval(node_ip)
|
|
print(node_ip)
|
|
# insert node ip to database
|
|
with sqlite3.connect("client.db") as db:
|
|
db.executemany(
|
|
"""
|
|
INSERT INTO node
|
|
(nodeip)
|
|
VALUES (?)
|
|
""",
|
|
[(ip,) for ip in node_ip],
|
|
)
|
|
db.commit()
|
|
print("Success add node ip")
|
|
else:
|
|
print("Failed:", response.status_code, response.text)
|
|
|
|
|
|
# send pk to others
|
|
@app.get("/get_pk")
|
|
async def get_pk():
|
|
global pk, sk
|
|
print(sk)
|
|
return {"pkx": str(pk[0]), "pky": str(pk[1])}
|
|
|
|
|
|
class pk_model(BaseModel):
|
|
pkx: str
|
|
pky: str
|
|
ip: str
|
|
|
|
|
|
# recieve pk from frontend
|
|
@app.post("/recieve_pk")
|
|
async def recieve_pk(pk: pk_model):
|
|
pkx = pk.pkx
|
|
pky = pk.pky
|
|
dest_ip = pk.ip
|
|
try:
|
|
threshold = 2
|
|
with sqlite3.connect("client.db") as db:
|
|
db.execute(
|
|
"""
|
|
INSERT INTO senderinfo
|
|
(ip, pkx, pky, threshold)
|
|
VALUES
|
|
(?, ?, ?, ?)
|
|
""",
|
|
(str(dest_ip), pkx, pky, threshold),
|
|
)
|
|
except Exception as e:
|
|
# raise error
|
|
print("Database error")
|
|
content = {"message": "Database Error", "error": str(e)}
|
|
return JSONResponse(content, status_code=400)
|
|
return {"message": "save pk in database"}
|
|
|
|
|
|
pk = (0, 0)
|
|
sk = 0
|
|
server_address = str
|
|
node_response = False
|
|
message = bytes
|
|
local_ip = get_own_ip()
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn # pylint: disable=e0401
|
|
|
|
uvicorn.run("client:app", host="0.0.0.0", port=8002, reload=True, log_level="debug")
|