From 1c41ab94b2f053b9baaf4a7227f4662127d3b6cb Mon Sep 17 00:00:00 2001
From: Smart-SangGe <2251250136@qq.com>
Date: Tue, 24 Oct 2023 22:40:12 +0800
Subject: [PATCH] fix: update

---
 src/client.py | 5 +++--
 src/node.py   | 6 +++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/client.py b/src/client.py
index cd71218..c718cb0 100644
--- a/src/client.py
+++ b/src/client.py
@@ -197,14 +197,15 @@ async def send_messages(
     # 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][0] + ":8001" + "/user_src"
         payload = {
             "source_ip": local_ip,
             "dest_ip": dest_ip,
-            "capsule_ct": base64.b64encode(str(capsule_ct)),
+            "capsule_ct": capsule_ct,
             "rk": rk_list[i],
         }
         print(payload)
diff --git a/src/node.py b/src/node.py
index 2cbf6b4..228e4a9 100644
--- a/src/node.py
+++ b/src/node.py
@@ -76,7 +76,7 @@ async def send_heartbeat_internal() -> None:
 class Req(BaseModel):
     source_ip: str
     dest_ip: str
-    capsule_ct: Tuple[capsule, bytes]
+    capsule_ct: Tuple[capsule, int]
     rk: Any
 
 
@@ -94,8 +94,8 @@ async def user_src(message: Req):
     """
     source_ip = message.source_ip
     dest_ip = message.dest_ip
-    b64capsule_ct = message.capsule_ct
-    capsule_ct = base64.b64decode(b64capsule_ct)
+    capsule, ct = message.capsule_ct
+    capsule_ct = (capsule, ct.to_bytes())
     rk = message.rk
 
     processed_message = ReEncrypt(rk, capsule_ct)