From 3467118d6c692eeb78a814bb9c7c8e175817fbb3 Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Tue, 17 Oct 2023 21:40:15 +0800
Subject: [PATCH 01/14] fix: remove setup function

---
 src/tpre.py | 21 ---------------------
 1 file changed, 21 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index a4eacd5..5f19e41 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -126,27 +126,6 @@ def jacobianMultiply(
 # 生成元    
 U = multiply(g, random.randint(0, sm2p256v1.P))
 
-# def Setup(sec: int) -> Tuple[CurveFp, Tuple[int, int], 
-#                              Tuple[int, int]]:
-#     '''
-#     params:
-#     sec: an init safety param
-    
-#     return:
-#     G: sm2 curve
-#     g: generator
-#     U: another generator
-#     '''
-    
-#     G = sm2p256v1
-    
-#     g = (sm2p256v1.Gx, sm2p256v1.Gy)
-    
-#     tmp_u = random.randint(0, sm2p256v1.P)
-#     U = multiply(g, tmp_u)
-    
-#     return G, g, U
-
 def hash2(double_G: Tuple[Tuple[int, int], Tuple[int, int]]) -> int:
     sm3 = Sm3() #pylint: disable=e0602
     for i in double_G:

From df16e661f2cfc12300a2474af8f3f2b62b62cb0c Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Tue, 17 Oct 2023 21:51:04 +0800
Subject: [PATCH 02/14] fix: rename hash

---
 src/tpre.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index 5f19e41..8a88f20 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -234,7 +234,7 @@ def Decrypt(sk_A: int,C:Tuple[Tuple[
     return dec_Data
 
 # GenerateRekey
-def H5(id: int, D: int) -> int:
+def hash5(id: int, D: int) -> int:
     sm3 = Sm3() #pylint: disable=e0602
     sm3.update(id.to_bytes(32))
     sm3.update(D.to_bytes(32))
@@ -242,7 +242,7 @@ def H5(id: int, D: int) -> int:
     hash = int.from_bytes(hash,'big') % G.P
     return hash
 
-def H6(triple_G: Tuple[Tuple[int, int], 
+def hash6(triple_G: Tuple[Tuple[int, int], 
                               Tuple[int, int],
                               Tuple[int, int]]) -> int:
         sm3 = Sm3() #pylint: disable=e0602
@@ -283,14 +283,14 @@ def GenerateReKey(sk_A, pk_B, N: int, T: int) -> list:
         f_modulus.append(random.randint(0, G.P - 1))
 
     # 计算D
-    D = H6((X_A, pk_B, multiply(pk_B, sk_A)))
+    D = hash6((X_A, pk_B, multiply(pk_B, sk_A)))
 
     # 计算KF
     KF = []
     for i in range(N):
         y = random.randint(0, G.P - 1)
         Y = multiply(g, y)
-        s_x = H5(i, D)         # id需要设置
+        s_x = hash5(i, D)         # id需要设置
         r_k = f(s_x, f_modulus, T)
         U1 = multiply(U, r_k)   
         kFrag = (i, r_k, X_A, U1)
@@ -379,10 +379,10 @@ def DecapsulateFrags(sk_B:int,pk_A:Tuple[int,int],cFrags:Tuple[Tuple[Tuple[int,i
         t = t+1        # 总共有t个片段，t为阈值
   
     pkab = multiply(pk_A,sk_B)     # pka^b
-    D = H6((pk_A,pk_B,pkab))
+    D = hash6((pk_A,pk_B,pkab))
     Sx = []
     for id in idlist:        #  从1到t
-        sxi = H5(id,D)       #  id 节点的编号
+        sxi = hash5(id,D)       #  id 节点的编号
         Sx.append(sxi)    
     bis= []    #  b ==> λ
     j = 1

From dff57fc3c6b0d071fef1b453b800c2e98dc6e3c0 Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Tue, 17 Oct 2023 21:54:52 +0800
Subject: [PATCH 03/14] fix: fix name error

---
 src/tpre.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index 8a88f20..4374bdb 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -403,8 +403,8 @@ def DecapsulateFrags(sk_B:int,pk_A:Tuple[int,int],cFrags:Tuple[Tuple[Tuple[int,i
         Vk = multiply(Vlist[k],bis[k])
         E2 = add(Ek,E2)   
         V2 = add(Vk,V2)
-    X_Ab = multiply(Xalist[0],b)     # X_A^b   X_A 的值是随机生成的xa，通过椭圆曲线上的倍点运算生成的固定的值
-    d = hash3((Xalist[0],pk_B,X_Ab))
+    X_Ab = multiply(X_Alist[0],sk_B)     # X_A^b   X_A 的值是随机生成的xa，通过椭圆曲线上的倍点运算生成的固定的值
+    d = hash3((X_Alist[0],pk_B,X_Ab))
     EV = add(E2,V2)    # E2 + V2
     EVd = multiply(EV,d)     # (E2 + V2)^d
     K = KDF(EVd)

From 5278a1fc64256a339e7fa5e6c6d36b90f0f3f03d Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Tue, 17 Oct 2023 22:04:49 +0800
Subject: [PATCH 04/14] fix: remove pk sk

---
 src/tpre.py | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index 4374bdb..739847f 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -193,8 +193,8 @@ def GenerateKeyPair(
     return public_key, secret_key
 
 # 生成A和B的公钥和私钥
-pk_A, sk_A = GenerateKeyPair(0, ())
-pk_B, sk_B = GenerateKeyPair(0, ())
+#pk_A, sk_A = GenerateKeyPair(0, ())
+#pk_B, sk_B = GenerateKeyPair(0, ())
 
 def Encrypt(pk: Tuple[int, int], m: int) -> Tuple[Tuple[
     Tuple[int, int],Tuple[int, int], int], int]:
@@ -254,6 +254,9 @@ def hash6(triple_G: Tuple[Tuple[int, int],
         return hash
 
 def f(x: int, f_modulus: list, T: int) -> int:
+    '''
+    
+    '''
     res = 0
     for i in range(T):
         res += f_modulus[i] * pow(x, i)
@@ -346,7 +349,8 @@ def ReEncrypt(kFrag:list,
 
 
 # N 是加密节点的数量，t是阈值
-def mergecfrag(N:int,t:int)->tuple[Tuple[Tuple[int,int],Tuple[int,int]
+def mergecfrag(sk_A: int, pk_A: Tuple[int, int], pk_B: Tuple[int, int],
+               N: int, t: int)->tuple[Tuple[Tuple[int,int],Tuple[int,int]
                               ,int,Tuple[int,int]], ...]:
     cfrags = ()
     kfrags = GenerateReKey(sk_A,pk_B,N,t)
@@ -360,8 +364,11 @@ def mergecfrag(N:int,t:int)->tuple[Tuple[Tuple[int,int],Tuple[int,int]
 
     
 
-def DecapsulateFrags(sk_B:int,pk_A:Tuple[int,int],cFrags:Tuple[Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]]]
-                  ,capsule:Tuple[Tuple[int,int],Tuple[int,int],int]) -> int:
+def DecapsulateFrags(sk_B:int, 
+                     pk_B: Tuple[int, int], 
+                     pk_A:Tuple[int,int],
+                     cFrags:Tuple[Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]]]
+                    ) -> int:
     '''
     return:
     K: sm4 key
@@ -413,13 +420,14 @@ def DecapsulateFrags(sk_B:int,pk_A:Tuple[int,int],cFrags:Tuple[Tuple[Tuple[int,i
 
 #  M = IAEAM(K,enc_Data)
 
-def DecryptFrags(sk_B:int,
-                 pk_A:Tuple[int,int],
-                 cFrags:Tuple[Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]]],
-                 C:Tuple[Tuple[Tuple[int,int],Tuple[int,int],int],int]
-                 )->int:
+def DecryptFrags(sk_B: int,
+                 pk_B: Tuple[int, int],
+                 pk_A: Tuple[int,int],
+                 cFrags: Tuple[Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]]],
+                 C: Tuple[Tuple[Tuple[int,int],Tuple[int,int],int],int]
+                 ) -> int:
     capsule,enc_Data = C   # 加密后的密文
-    K = DecapsulateFrags(sk_B,pk_A,cFrags,capsule)
+    K = DecapsulateFrags(sk_B, pk_B, pk_A,cFrags)
     
     iv = b'tpretpretpretpre'
     sm4_dec = Sm4Cbc(K, iv, DO_DECRYPT) #pylint: disable= e0602

From f9e22fa39ce8e90a70d0f613fe179304d031aa4d Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Tue, 17 Oct 2023 22:06:04 +0800
Subject: [PATCH 05/14] feat: start trace demo

---
 .gitignore  | 1 -
 src/demo.py | 3 +++
 2 files changed, 3 insertions(+), 1 deletion(-)
 create mode 100644 src/demo.py

diff --git a/.gitignore b/.gitignore
index c8705f0..80d5f5f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,3 @@ test.py
 example.py 
 
 ReEncrypt.py
-src/demo.py
diff --git a/src/demo.py b/src/demo.py
new file mode 100644
index 0000000..ad5f059
--- /dev/null
+++ b/src/demo.py
@@ -0,0 +1,3 @@
+from tpre import *
+
+

From 9e8760e5d9e372a74695547a0e825ad4495ba5ca Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Tue, 17 Oct 2023 22:40:17 +0800
Subject: [PATCH 06/14] feat: add point type and capsule type

---
 src/tpre.py | 73 ++++++++++++++++++++++++++++-------------------------
 1 file changed, 38 insertions(+), 35 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index 739847f..d512af9 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -2,6 +2,9 @@ from gmssl import * #pylint: disable = e0401
 from typing import Tuple, Callable
 import random
 
+point = Tuple[int, int]
+capsule = Tuple[point, point, int]
+
 # 生成密钥对模块
 class CurveFp:
     def __init__(self, A, B, P, N, Gx, Gy, name):
@@ -29,13 +32,13 @@ G = sm2p256v1
 # 生成元
 g = (sm2p256v1.Gx, sm2p256v1.Gy)
 
-def multiply(a: Tuple[int, int], n: int) -> Tuple[int, int]:
+def multiply(a: point, n: int) -> point:
     N = sm2p256v1.N
     A = sm2p256v1.A
     P = sm2p256v1.P 
     return fromJacobian(jacobianMultiply(toJacobian(a), n, N, A, P), P)
 
-def add(a: Tuple[int, int], b: Tuple[int, int]) -> Tuple[int, int]:
+def add(a: point, b: point) -> point:
     A = sm2p256v1.A
     P = sm2p256v1.P
     return fromJacobian(jacobianAdd(toJacobian(a), toJacobian(b), A, P), P)
@@ -51,11 +54,11 @@ def inv(a: int, n: int) -> int:
         lm, low, hm, high = nm, new, lm, low
     return lm % n
  
-def toJacobian(Xp_Yp: Tuple[int, int]) -> Tuple[int, int, int]:
+def toJacobian(Xp_Yp: point) -> Tuple[int, int, int]:
     Xp, Yp = Xp_Yp
     return (Xp, Yp, 1)
 
-def fromJacobian(Xp_Yp_Zp: Tuple[int, int, int], P: int) -> Tuple[int, int]:
+def fromJacobian(Xp_Yp_Zp: Tuple[int, int, int], P: int) -> point:
     Xp, Yp, Zp = Xp_Yp_Zp
     z = inv(Zp, P)
     return ((Xp * z ** 2) % P, (Yp * z ** 3) % P)
@@ -126,7 +129,7 @@ def jacobianMultiply(
 # 生成元    
 U = multiply(g, random.randint(0, sm2p256v1.P))
 
-def hash2(double_G: Tuple[Tuple[int, int], Tuple[int, int]]) -> int:
+def hash2(double_G: Tuple[point, point]) -> int:
     sm3 = Sm3() #pylint: disable=e0602
     for i in double_G:
         for j in i:
@@ -135,9 +138,9 @@ def hash2(double_G: Tuple[Tuple[int, int], Tuple[int, int]]) -> int:
     digest = int.from_bytes(digest,'big') % sm2p256v1.P
     return digest
 
-def hash3(triple_G: Tuple[Tuple[int, int], 
-                            Tuple[int, int],
-                            Tuple[int, int]]) -> int:
+def hash3(triple_G: Tuple[point, 
+                            point,
+                            point]) -> int:
     sm3 = Sm3() #pylint: disable=e0602
     for i in triple_G:
         for j in i:
@@ -146,9 +149,9 @@ def hash3(triple_G: Tuple[Tuple[int, int],
     digest = int.from_bytes(digest, 'big') % sm2p256v1.P
     return digest
 
-def hash4(triple_G: Tuple[Tuple[int, int],
-                            Tuple[int, int],
-                            Tuple[int, int]],
+def hash4(triple_G: Tuple[point,
+                            point,
+                            point],
             Zp: int) -> int:
     sm3 = Sm3() #pylint: disable=e0602
     for i in triple_G:
@@ -159,7 +162,7 @@ def hash4(triple_G: Tuple[Tuple[int, int],
     digest = int.from_bytes(digest, 'big') % sm2p256v1.P
     return digest
 
-def KDF(G: Tuple[int, int]) -> int:
+def KDF(G: point) -> int:
     sm3 = Sm3() #pylint: disable=e0602
     for i in G:
         sm3.update(i.to_bytes(32))
@@ -171,7 +174,7 @@ def KDF(G: Tuple[int, int]) -> int:
 def GenerateKeyPair(
     lamda_parma: int, 
     public_params: tuple
-    ) -> Tuple[Tuple[int, int], int]:
+    ) -> Tuple[point, int]:
     '''
     params: 
     lamda_param: an init safety param 
@@ -196,8 +199,8 @@ def GenerateKeyPair(
 #pk_A, sk_A = GenerateKeyPair(0, ())
 #pk_B, sk_B = GenerateKeyPair(0, ())
 
-def Encrypt(pk: Tuple[int, int], m: int) -> Tuple[Tuple[
-    Tuple[int, int],Tuple[int, int], int], int]:
+def Encrypt(pk: point, m: int) -> Tuple[Tuple[
+    point,point, int], int]:
     enca = Encapsulate(pk)
     K = enca[0].to_bytes()
     capsule = enca[1]
@@ -211,7 +214,7 @@ def Encrypt(pk: Tuple[int, int], m: int) -> Tuple[Tuple[
     enc_message = (capsule, enc_Data)
     return enc_message
 
-def Decapsulate(ska:int,capsule:Tuple[Tuple[int,int],Tuple[int,int],int]) -> int:
+def Decapsulate(ska:int,capsule:capsule) -> int:
     E,V,s = capsule
     EVa=multiply(add(E,V), ska)    # (E*V)^ska
     K = KDF(EVa)
@@ -219,7 +222,7 @@ def Decapsulate(ska:int,capsule:Tuple[Tuple[int,int],Tuple[int,int],int]) -> int
     return K
 
 def Decrypt(sk_A: int,C:Tuple[Tuple[
-    Tuple[int, int],Tuple[int, int], int], int]) ->int:
+    point, point, int], int]) ->int:
     '''
     params:
     sk_A: secret key
@@ -242,9 +245,9 @@ def hash5(id: int, D: int) -> int:
     hash = int.from_bytes(hash,'big') % G.P
     return hash
 
-def hash6(triple_G: Tuple[Tuple[int, int], 
-                              Tuple[int, int],
-                              Tuple[int, int]]) -> int:
+def hash6(triple_G: Tuple[point, 
+                              point,
+                              point]) -> int:
         sm3 = Sm3() #pylint: disable=e0602
         for i in triple_G:
             for j in i:
@@ -262,7 +265,7 @@ def f(x: int, f_modulus: list, T: int) -> int:
         res += f_modulus[i] * pow(x, i)
     return res
 
-def GenerateReKey(sk_A, pk_B, N: int, T: int) -> list:
+def GenerateReKey(sk_A: int, pk_B: point, N: int, T: int) -> list:
     '''
     param: 
     skA, pkB, N(节点总数), T(阈值)
@@ -301,7 +304,7 @@ def GenerateReKey(sk_A, pk_B, N: int, T: int) -> list:
 
     return KF
 
-def Encapsulate(pk_A: Tuple[int, int]) -> Tuple[int, Tuple[Tuple[int, int], Tuple[int, int], int]]:
+def Encapsulate(pk_A: point) -> Tuple[int, capsule]:
     r = random.randint(0, G.P - 1)
     u = random.randint(0, G.P - 1)
     E = multiply(g, r)
@@ -312,7 +315,7 @@ def Encapsulate(pk_A: Tuple[int, int]) -> Tuple[int, Tuple[Tuple[int, int], Tupl
     capsule = (E, V, s)
     return (K, capsule)
 
-def Checkcapsule(capsule:Tuple[Tuple[int,int],Tuple[int,int],int]) -> bool:  # 验证胶囊的有效性
+def Checkcapsule(capsule:capsule) -> bool:  # 验证胶囊的有效性
     E,V,s = capsule
     h2 = hash2((E,V))
     g = (sm2p256v1.Gx, sm2p256v1.Gy)
@@ -326,7 +329,7 @@ def Checkcapsule(capsule:Tuple[Tuple[int,int],Tuple[int,int],int]) -> bool:  # 
         
     return flag 
 
-def ReEncapsulate(kFrag:list,capsule:Tuple[Tuple[int,int],Tuple[int,int],int]) -> Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]] :
+def ReEncapsulate(kFrag:list,capsule:capsule) -> Tuple[point,point,int,point] :
     id,rk,Xa,U1 = kFrag
     E,V,s = capsule
     if not Checkcapsule(capsule):
@@ -340,7 +343,7 @@ def ReEncapsulate(kFrag:list,capsule:Tuple[Tuple[int,int],Tuple[int,int],int]) -
     
     # 重加密函数
 def ReEncrypt(kFrag:list,
-              C:Tuple[Tuple[Tuple[int,int],Tuple[int,int],int],int])->Tuple[Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]],int] :
+              C:Tuple[capsule,int])->Tuple[Tuple[point,point,int,point],int] :
     capsule,enc_Data = C
 
     cFrag = ReEncapsulate(kFrag,capsule)
@@ -349,9 +352,9 @@ def ReEncrypt(kFrag:list,
 
 
 # N 是加密节点的数量，t是阈值
-def mergecfrag(sk_A: int, pk_A: Tuple[int, int], pk_B: Tuple[int, int],
-               N: int, t: int)->tuple[Tuple[Tuple[int,int],Tuple[int,int]
-                              ,int,Tuple[int,int]], ...]:
+def mergecfrag(sk_A: int, pk_A: point, pk_B: point,
+               N: int, t: int)->tuple[Tuple[point,point
+                              ,int,point], ...]:
     cfrags = ()
     kfrags = GenerateReKey(sk_A,pk_B,N,t)
     result  = Encapsulate(pk_A)
@@ -365,9 +368,9 @@ def mergecfrag(sk_A: int, pk_A: Tuple[int, int], pk_B: Tuple[int, int],
     
 
 def DecapsulateFrags(sk_B:int, 
-                     pk_B: Tuple[int, int], 
-                     pk_A:Tuple[int,int],
-                     cFrags:Tuple[Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]]]
+                     pk_B: point, 
+                     pk_A:point,
+                     cFrags:Tuple[Tuple[point,point,int,point]]
                     ) -> int:
     '''
     return:
@@ -421,10 +424,10 @@ def DecapsulateFrags(sk_B:int,
 #  M = IAEAM(K,enc_Data)
 
 def DecryptFrags(sk_B: int,
-                 pk_B: Tuple[int, int],
-                 pk_A: Tuple[int,int],
-                 cFrags: Tuple[Tuple[Tuple[int,int],Tuple[int,int],int,Tuple[int,int]]],
-                 C: Tuple[Tuple[Tuple[int,int],Tuple[int,int],int],int]
+                 pk_B: point,
+                 pk_A: point,
+                 cFrags: Tuple[Tuple[point,point,int,point]],
+                 C: Tuple[capsule,int]
                  ) -> int:
     capsule,enc_Data = C   # 加密后的密文
     K = DecapsulateFrags(sk_B, pk_B, pk_A,cFrags)

From c0494494a75f28861c17efc1f5ae5b050f36b98f Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Tue, 17 Oct 2023 23:00:28 +0800
Subject: [PATCH 07/14] feat: finish demo

---
 src/demo.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/demo.py b/src/demo.py
index ad5f059..f511597 100644
--- a/src/demo.py
+++ b/src/demo.py
@@ -1,3 +1,31 @@
 from tpre import *
 
+# 1
+pk_a, sk_a = GenerateKeyPair(1, ())
+m = b'hello world'
+m = int.from_bytes(m)
+
+# 2
+capsule_ct = Encrypt(pk_a, m)
+
+# 3
+pk_b, sk_b = GenerateKeyPair(1, ())
+
+N = 20
+T = 10
+
+# 5
+rekeys =  GenerateReKey(sk_a, pk_b, N, T)
+
+# 7
+cfrag_cts = []
+
+for rekey in rekeys:
+    cfrag_ct = ReEncrypt(rekey, capsule_ct)
+    cfrag_cts.append(cfrag_ct)
+    
+# 9
+cfrags = mergecfrag(cfrag_cts)
+m = DecryptFrags(sk_b, pk_b, pk_a, cfrags)
+
 

From c2698455de65e5d66dca7041171530bb1fbe6931 Mon Sep 17 00:00:00 2001
From: ccyj <2384899431@qq.com>
Date: Wed, 18 Oct 2023 10:16:49 +0800
Subject: [PATCH 08/14] =?UTF-8?q?fix:=20=E4=BF=AE=E6=94=B9=E4=BA=86merge?=
 =?UTF-8?q?=20=E5=92=8CDecryptfrags=20=E5=87=BD=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/tpre.py | 33 ++++++++++++++++-----------------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index d512af9..b9e242b 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -351,18 +351,16 @@ def ReEncrypt(kFrag:list,
 # capsule, enc_Data = C
 
 
-# N 是加密节点的数量，t是阈值
-def mergecfrag(sk_A: int, pk_A: point, pk_B: point,
-               N: int, t: int)->tuple[Tuple[point,point
-                              ,int,point], ...]:
-    cfrags = ()
-    kfrags = GenerateReKey(sk_A,pk_B,N,t)
-    result  = Encapsulate(pk_A)
-    K,capsule = result  
-    for kfrag in kfrags:
-        cfrag = ReEncapsulate(kfrag,capsule)
-        cfrags = cfrags + (cfrag,)
-    
+# 将加密节点加密后产生的t个（capsule,ct）合并在一起，产生cfrags = {{capsule1,capsule2,...},ct}
+def mergecfrag(cfrag_cts:list)->list:
+    ct_list = []
+    cfrags_list = []
+    cfrags = []
+    for cfrag_ct in cfrag_cts:
+        cfrags_list.append(cfrag_ct[0])
+        cfrags_list.append(cfrag_ct[1])
+    cfrags.append(cfrags_list)
+    cfrags.append(ct_list[0])
     return cfrags
 
     
@@ -370,12 +368,13 @@ def mergecfrag(sk_A: int, pk_A: point, pk_B: point,
 def DecapsulateFrags(sk_B:int, 
                      pk_B: point, 
                      pk_A:point,
-                     cFrags:Tuple[Tuple[point,point,int,point]]
+                     cFrags:list
                     ) -> int:
     '''
     return:
     K: sm4 key
     '''
+
     Elist = []
     Vlist = []
     idlist = []
@@ -423,14 +422,14 @@ def DecapsulateFrags(sk_B:int,
 
 #  M = IAEAM(K,enc_Data)
 
+# cfrags = {{capsule1,capsule2,...},ct} ,ct->en_Data
 def DecryptFrags(sk_B: int,
                  pk_B: point,
                  pk_A: point,
-                 cFrags: Tuple[Tuple[point,point,int,point]],
-                 C: Tuple[capsule,int]
+                 cfrags:list
                  ) -> int:
-    capsule,enc_Data = C   # 加密后的密文
-    K = DecapsulateFrags(sk_B, pk_B, pk_A,cFrags)
+    capsules,enc_Data = cfrags   # 加密后的密文
+    K = DecapsulateFrags(sk_B, pk_B, pk_A, capsules)
     
     iv = b'tpretpretpretpre'
     sm4_dec = Sm4Cbc(K, iv, DO_DECRYPT) #pylint: disable= e0602

From c7a942634ecc3ede50136cf1140845440161f7f7 Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Wed, 18 Oct 2023 11:00:32 +0800
Subject: [PATCH 09/14] fix: fix bugs

---
 src/tpre.py | 32 ++++++++++++++++++++------------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index b9e242b..2c3d618 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -164,11 +164,15 @@ def hash4(triple_G: Tuple[point,
 
 def KDF(G: point) -> int:
     sm3 = Sm3() #pylint: disable=e0602
+    print(G)
     for i in G:
         sm3.update(i.to_bytes(32))
-    digest = sm3.digest(32)
+    digest = sm3.digest()
     digest = digest
     digest = int.from_bytes(digest, 'big') % sm2p256v1.P
+    mask_128bit = (1 << 128) - 1
+    digest = digest & mask_128bit
+    print('key =',digest)
     return digest
 
 def GenerateKeyPair(
@@ -199,17 +203,16 @@ def GenerateKeyPair(
 #pk_A, sk_A = GenerateKeyPair(0, ())
 #pk_B, sk_B = GenerateKeyPair(0, ())
 
-def Encrypt(pk: point, m: int) -> Tuple[Tuple[
-    point,point, int], int]:
+def Encrypt(pk: point, m: bytes) -> Tuple[Tuple[
+    point,point, int], bytes]:
     enca = Encapsulate(pk)
-    K = enca[0].to_bytes()
+    K = enca[0].to_bytes(16)
     capsule = enca[1]
     if len(K) != 16:
         raise ValueError("invalid key length")
     iv = b'tpretpretpretpre'
     sm4_enc = Sm4Cbc(K, iv, DO_ENCRYPT) #pylint: disable=e0602
-    plain_Data = m.to_bytes(32)
-    enc_Data = sm4_enc.update(plain_Data)
+    enc_Data = sm4_enc.update(m)
     enc_Data += sm4_enc.finish()
     enc_message = (capsule, enc_Data)
     return enc_message
@@ -222,7 +225,7 @@ def Decapsulate(ska:int,capsule:capsule) -> int:
     return K
 
 def Decrypt(sk_A: int,C:Tuple[Tuple[
-    point, point, int], int]) ->int:
+    point, point, int], bytes]) ->int:
     '''
     params:
     sk_A: secret key
@@ -358,7 +361,7 @@ def mergecfrag(cfrag_cts:list)->list:
     cfrags = []
     for cfrag_ct in cfrag_cts:
         cfrags_list.append(cfrag_ct[0])
-        cfrags_list.append(cfrag_ct[1])
+        ct_list.append(cfrag_ct[1])
     cfrags.append(cfrags_list)
     cfrags.append(ct_list[0])
     return cfrags
@@ -427,12 +430,17 @@ def DecryptFrags(sk_B: int,
                  pk_B: point,
                  pk_A: point,
                  cfrags:list
-                 ) -> int:
+                 ) -> bytes:
     capsules,enc_Data = cfrags   # 加密后的密文
     K = DecapsulateFrags(sk_B, pk_B, pk_A, capsules)
-    
+    K = K.to_bytes(16)
     iv = b'tpretpretpretpre'
     sm4_dec = Sm4Cbc(K, iv, DO_DECRYPT) #pylint: disable= e0602
-    dec_Data = sm4_dec.update(enc_Data)
-    dec_Data += sm4_dec.finish()
+    try:
+        dec_Data = sm4_dec.update(enc_Data)
+        dec_Data += sm4_dec.finish()
+    except Exception as e:
+        print(e)
+        print("key error")
+        dec_Data = b''
     return dec_Data
\ No newline at end of file

From 1958f95c24d51b79767457448900d3c051fcde98 Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Wed, 18 Oct 2023 15:22:00 +0800
Subject: [PATCH 10/14] feat: add debug trace

---
 src/tpre.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/tpre.py b/src/tpre.py
index 2c3d618..54c086a 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -1,6 +1,7 @@
 from gmssl import * #pylint: disable = e0401 
 from typing import Tuple, Callable
 import random
+import traceback
 
 point = Tuple[int, int]
 capsule = Tuple[point, point, int]
@@ -173,6 +174,7 @@ def KDF(G: point) -> int:
     mask_128bit = (1 << 128) - 1
     digest = digest & mask_128bit
     print('key =',digest)
+    traceback.print_stack()
     return digest
 
 def GenerateKeyPair(
@@ -403,7 +405,8 @@ def DecapsulateFrags(sk_B:int,
     for i in range(t):
         for j in range(t):
             if j == i:
-                j=j+1
+                # j=j+1
+                continue
             else:
                 bi = bi * (Sx[j]//(Sx[j]-Sx[i]))    # 暂定整除
         bis.append(bi)

From 02fbfbbd7034a2e68f250259b74614639423c671 Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Wed, 18 Oct 2023 15:26:52 +0800
Subject: [PATCH 11/14] style: use black to modify style

---
 src/tpre.py | 359 +++++++++++++++++++++++++++-------------------------
 1 file changed, 184 insertions(+), 175 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index 54c086a..a8cb0ca 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -1,4 +1,4 @@
-from gmssl import * #pylint: disable = e0401 
+from gmssl import *  # pylint: disable = e0401
 from typing import Tuple, Callable
 import random
 import traceback
@@ -6,6 +6,7 @@ import traceback
 point = Tuple[int, int]
 capsule = Tuple[point, point, int]
 
+
 # 生成密钥对模块
 class CurveFp:
     def __init__(self, A, B, P, N, Gx, Gy, name):
@@ -16,7 +17,8 @@ class CurveFp:
         self.Gx = Gx
         self.Gy = Gy
         self.name = name
-        
+
+
 sm2p256v1 = CurveFp(
     name="sm2p256v1",
     A=0xFFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFC,
@@ -24,7 +26,7 @@ sm2p256v1 = CurveFp(
     P=0xFFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFF,
     N=0xFFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFF7203DF6B21C6052B53BBF40939D54123,
     Gx=0x32C4AE2C1F1981195F9904466A39C9948FE30BBFF2660BE1715A4589334C74C7,
-    Gy=0xBC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0
+    Gy=0xBC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0,
 )
 
 # 椭圆曲线
@@ -33,17 +35,20 @@ G = sm2p256v1
 # 生成元
 g = (sm2p256v1.Gx, sm2p256v1.Gy)
 
+
 def multiply(a: point, n: int) -> point:
     N = sm2p256v1.N
     A = sm2p256v1.A
-    P = sm2p256v1.P 
+    P = sm2p256v1.P
     return fromJacobian(jacobianMultiply(toJacobian(a), n, N, A, P), P)
 
+
 def add(a: point, b: point) -> point:
     A = sm2p256v1.A
     P = sm2p256v1.P
     return fromJacobian(jacobianAdd(toJacobian(a), toJacobian(b), A, P), P)
- 
+
+
 def inv(a: int, n: int) -> int:
     if a == 0:
         return 0
@@ -54,33 +59,36 @@ def inv(a: int, n: int) -> int:
         nm, new = hm - lm * r, high - low * r
         lm, low, hm, high = nm, new, lm, low
     return lm % n
- 
+
+
 def toJacobian(Xp_Yp: point) -> Tuple[int, int, int]:
     Xp, Yp = Xp_Yp
     return (Xp, Yp, 1)
 
+
 def fromJacobian(Xp_Yp_Zp: Tuple[int, int, int], P: int) -> point:
     Xp, Yp, Zp = Xp_Yp_Zp
     z = inv(Zp, P)
-    return ((Xp * z ** 2) % P, (Yp * z ** 3) % P)
- 
-def jacobianDouble(Xp_Yp_Zp: Tuple[int, int, int], A: int, P: int) -> Tuple[int, int, int]:
+    return ((Xp * z**2) % P, (Yp * z**3) % P)
+
+
+def jacobianDouble(
+    Xp_Yp_Zp: Tuple[int, int, int], A: int, P: int
+) -> Tuple[int, int, int]:
     Xp, Yp, Zp = Xp_Yp_Zp
     if not Yp:
         return (0, 0, 0)
-    ysq = (Yp ** 2) % P
+    ysq = (Yp**2) % P
     S = (4 * Xp * ysq) % P
-    M = (3 * Xp ** 2 + A * Zp ** 4) % P
-    nx = (M ** 2 - 2 * S) % P
-    ny = (M * (S - nx) - 8 * ysq ** 2) % P
+    M = (3 * Xp**2 + A * Zp**4) % P
+    nx = (M**2 - 2 * S) % P
+    ny = (M * (S - nx) - 8 * ysq**2) % P
     nz = (2 * Yp * Zp) % P
     return (nx, ny, nz)
- 
+
+
 def jacobianAdd(
-    Xp_Yp_Zp: Tuple[int, int, int], 
-    Xq_Yq_Zq: Tuple[int, int, int], 
-    A: int, 
-    P: int
+    Xp_Yp_Zp: Tuple[int, int, int], Xq_Yq_Zq: Tuple[int, int, int], A: int, P: int
 ) -> Tuple[int, int, int]:
     Xp, Yp, Zp = Xp_Yp_Zp
     Xq, Yq, Zq = Xq_Yq_Zq
@@ -88,10 +96,10 @@ def jacobianAdd(
         return (Xq, Yq, Zq)
     if not Yq:
         return (Xp, Yp, Zp)
-    U1 = (Xp * Zq ** 2) % P
-    U2 = (Xq * Zp ** 2) % P
-    S1 = (Yp * Zq ** 3) % P
-    S2 = (Yq * Zp ** 3) % P
+    U1 = (Xp * Zq**2) % P
+    U2 = (Xq * Zp**2) % P
+    S1 = (Yp * Zq**3) % P
+    S2 = (Yq * Zp**3) % P
     if U1 == U2:
         if S1 != S2:
             return (0, 0, 1)
@@ -101,19 +109,15 @@ def jacobianAdd(
     H2 = (H * H) % P
     H3 = (H * H2) % P
     U1H2 = (U1 * H2) % P
-    nx = (R ** 2 - H3 - 2 * U1H2) % P
+    nx = (R**2 - H3 - 2 * U1H2) % P
     ny = (R * (U1H2 - nx) - S1 * H3) % P
     nz = (H * Zp * Zq) % P
     return (nx, ny, nz)
- 
-def jacobianMultiply(
-    Xp_Yp_Zp: Tuple[int, int, int], 
-    n: int, 
-    N: int, 
-    A: int, 
-    P: int
-    ) -> Tuple[int, int, int]:
 
+
+def jacobianMultiply(
+    Xp_Yp_Zp: Tuple[int, int, int], n: int, N: int, A: int, P: int
+) -> Tuple[int, int, int]:
     Xp, Yp, Zp = Xp_Yp_Zp
     if Yp == 0 or n == 0:
         return (0, 0, 1)
@@ -124,166 +128,170 @@ def jacobianMultiply(
     if (n % 2) == 0:
         return jacobianDouble(jacobianMultiply((Xp, Yp, Zp), n // 2, N, A, P), A, P)
     if (n % 2) == 1:
-        return jacobianAdd(jacobianDouble(jacobianMultiply((Xp, Yp, Zp), n // 2, N, A, P), A, P), (Xp, Yp, Zp), A, P)
+        return jacobianAdd(
+            jacobianDouble(jacobianMultiply((Xp, Yp, Zp), n // 2, N, A, P), A, P),
+            (Xp, Yp, Zp),
+            A,
+            P,
+        )
     raise ValueError("jacobian Multiply error")
 
-# 生成元    
+
+# 生成元
 U = multiply(g, random.randint(0, sm2p256v1.P))
 
+
 def hash2(double_G: Tuple[point, point]) -> int:
-    sm3 = Sm3() #pylint: disable=e0602
+    sm3 = Sm3()  # pylint: disable=e0602
     for i in double_G:
         for j in i:
             sm3.update(j.to_bytes(32))
     digest = sm3.digest()
-    digest = int.from_bytes(digest,'big') % sm2p256v1.P
+    digest = int.from_bytes(digest, "big") % sm2p256v1.P
     return digest
 
-def hash3(triple_G: Tuple[point, 
-                            point,
-                            point]) -> int:
-    sm3 = Sm3() #pylint: disable=e0602
+
+def hash3(triple_G: Tuple[point, point, point]) -> int:
+    sm3 = Sm3()  # pylint: disable=e0602
     for i in triple_G:
         for j in i:
             sm3.update(j.to_bytes(32))
     digest = sm3.digest()
-    digest = int.from_bytes(digest, 'big') % sm2p256v1.P
+    digest = int.from_bytes(digest, "big") % sm2p256v1.P
     return digest
 
-def hash4(triple_G: Tuple[point,
-                            point,
-                            point],
-            Zp: int) -> int:
-    sm3 = Sm3() #pylint: disable=e0602
+
+def hash4(triple_G: Tuple[point, point, point], Zp: int) -> int:
+    sm3 = Sm3()  # pylint: disable=e0602
     for i in triple_G:
         for j in i:
             sm3.update(j.to_bytes(32))
     sm3.update(Zp.to_bytes(32))
     digest = sm3.digest()
-    digest = int.from_bytes(digest, 'big') % sm2p256v1.P
+    digest = int.from_bytes(digest, "big") % sm2p256v1.P
     return digest
 
+
 def KDF(G: point) -> int:
-    sm3 = Sm3() #pylint: disable=e0602
+    sm3 = Sm3()  # pylint: disable=e0602
     print(G)
     for i in G:
         sm3.update(i.to_bytes(32))
     digest = sm3.digest()
     digest = digest
-    digest = int.from_bytes(digest, 'big') % sm2p256v1.P
+    digest = int.from_bytes(digest, "big") % sm2p256v1.P
     mask_128bit = (1 << 128) - 1
     digest = digest & mask_128bit
-    print('key =',digest)
+    print("key =", digest)
     traceback.print_stack()
     return digest
 
-def GenerateKeyPair(
-    lamda_parma: int, 
-    public_params: tuple
-    ) -> Tuple[point, int]:
-    '''
-    params: 
-    lamda_param: an init safety param 
+
+def GenerateKeyPair(lamda_parma: int, public_params: tuple) -> Tuple[point, int]:
+    """
+    params:
+    lamda_param: an init safety param
     public_params: curve params
-    
+
     return:
     public_key, secret_key
-    '''
-    sm2 = Sm2Key() #pylint: disable=e0602
+    """
+    sm2 = Sm2Key()  # pylint: disable=e0602
     sm2.generate_key()
-    
-    public_key_x = int.from_bytes(bytes(sm2.public_key.x),"big")
-    public_key_y = int.from_bytes(bytes(sm2.public_key.y),"big")
+
+    public_key_x = int.from_bytes(bytes(sm2.public_key.x), "big")
+    public_key_y = int.from_bytes(bytes(sm2.public_key.y), "big")
     public_key = (public_key_x, public_key_y)
-      
-    
-    secret_key = int.from_bytes(bytes(sm2.private_key),"big")
-    
+
+    secret_key = int.from_bytes(bytes(sm2.private_key), "big")
+
     return public_key, secret_key
 
-# 生成A和B的公钥和私钥
-#pk_A, sk_A = GenerateKeyPair(0, ())
-#pk_B, sk_B = GenerateKeyPair(0, ())
 
-def Encrypt(pk: point, m: bytes) -> Tuple[Tuple[
-    point,point, int], bytes]:
+# 生成A和B的公钥和私钥
+# pk_A, sk_A = GenerateKeyPair(0, ())
+# pk_B, sk_B = GenerateKeyPair(0, ())
+
+
+def Encrypt(pk: point, m: bytes) -> Tuple[Tuple[point, point, int], bytes]:
     enca = Encapsulate(pk)
     K = enca[0].to_bytes(16)
     capsule = enca[1]
     if len(K) != 16:
         raise ValueError("invalid key length")
-    iv = b'tpretpretpretpre'
-    sm4_enc = Sm4Cbc(K, iv, DO_ENCRYPT) #pylint: disable=e0602
+    iv = b"tpretpretpretpre"
+    sm4_enc = Sm4Cbc(K, iv, DO_ENCRYPT)  # pylint: disable=e0602
     enc_Data = sm4_enc.update(m)
     enc_Data += sm4_enc.finish()
     enc_message = (capsule, enc_Data)
     return enc_message
 
-def Decapsulate(ska:int,capsule:capsule) -> int:
-    E,V,s = capsule
-    EVa=multiply(add(E,V), ska)    # (E*V)^ska
+
+def Decapsulate(ska: int, capsule: capsule) -> int:
+    E, V, s = capsule
+    EVa = multiply(add(E, V), ska)  # (E*V)^ska
     K = KDF(EVa)
 
     return K
 
-def Decrypt(sk_A: int,C:Tuple[Tuple[
-    point, point, int], bytes]) ->int:
-    '''
+
+def Decrypt(sk_A: int, C: Tuple[Tuple[point, point, int], bytes]) -> int:
+    """
     params:
     sk_A: secret key
-    C: (capsule, enc_data) 
-    '''
-    capsule,enc_Data = C
-    K = Decapsulate(sk_A,capsule)
-    iv = b'tpretpretpretpre'
-    sm4_dec = Sm4Cbc(K, iv, DO_DECRYPT) #pylint: disable= e0602
+    C: (capsule, enc_data)
+    """
+    capsule, enc_Data = C
+    K = Decapsulate(sk_A, capsule)
+    iv = b"tpretpretpretpre"
+    sm4_dec = Sm4Cbc(K, iv, DO_DECRYPT)  # pylint: disable= e0602
     dec_Data = sm4_dec.update(enc_Data)
     dec_Data += sm4_dec.finish()
     return dec_Data
 
+
 # GenerateRekey
 def hash5(id: int, D: int) -> int:
-    sm3 = Sm3() #pylint: disable=e0602
+    sm3 = Sm3()  # pylint: disable=e0602
     sm3.update(id.to_bytes(32))
     sm3.update(D.to_bytes(32))
     hash = sm3.digest()
-    hash = int.from_bytes(hash,'big') % G.P
+    hash = int.from_bytes(hash, "big") % G.P
+    return hash
+
+
+def hash6(triple_G: Tuple[point, point, point]) -> int:
+    sm3 = Sm3()  # pylint: disable=e0602
+    for i in triple_G:
+        for j in i:
+            sm3.update(j.to_bytes(32))
+    hash = sm3.digest()
+    hash = int.from_bytes(hash, "big") % G.P
     return hash
 
-def hash6(triple_G: Tuple[point, 
-                              point,
-                              point]) -> int:
-        sm3 = Sm3() #pylint: disable=e0602
-        for i in triple_G:
-            for j in i:
-                sm3.update(j.to_bytes(32))
-        hash = sm3.digest()
-        hash = int.from_bytes(hash,'big') % G.P
-        return hash
 
 def f(x: int, f_modulus: list, T: int) -> int:
-    '''
-    
-    '''
+    """ """
     res = 0
     for i in range(T):
         res += f_modulus[i] * pow(x, i)
     return res
 
+
 def GenerateReKey(sk_A: int, pk_B: point, N: int, T: int) -> list:
-    '''
-    param: 
+    """
+    param:
     skA, pkB, N(节点总数), T(阈值)
-    return: 
+    return:
     rki(0 <= i <= N-1)
-    '''
+    """
     # 计算临时密钥对(x_A, X_A)
     x_A = random.randint(0, G.P - 1)
-    X_A = multiply(g, x_A)                
+    X_A = multiply(g, x_A)
 
     # d是Bob的密钥对与临时密钥对的非交互式Diffie-Hellman密钥交换的结果
-    d = hash3((X_A, pk_B, multiply(pk_B, x_A)))   
-    
+    d = hash3((X_A, pk_B, multiply(pk_B, x_A)))
+
     # 计算多项式系数, 确定代理节点的ID(一个点)
     f_modulus = []
     # 计算f0
@@ -301,14 +309,15 @@ def GenerateReKey(sk_A: int, pk_B: point, N: int, T: int) -> list:
     for i in range(N):
         y = random.randint(0, G.P - 1)
         Y = multiply(g, y)
-        s_x = hash5(i, D)         # id需要设置
+        s_x = hash5(i, D)  # id需要设置
         r_k = f(s_x, f_modulus, T)
-        U1 = multiply(U, r_k)   
+        U1 = multiply(U, r_k)
         kFrag = (i, r_k, X_A, U1)
         KF.append(kFrag)
 
     return KF
 
+
 def Encapsulate(pk_A: point) -> Tuple[int, capsule]:
     r = random.randint(0, G.P - 1)
     u = random.randint(0, G.P - 1)
@@ -320,44 +329,51 @@ def Encapsulate(pk_A: point) -> Tuple[int, capsule]:
     capsule = (E, V, s)
     return (K, capsule)
 
-def Checkcapsule(capsule:capsule) -> bool:  # 验证胶囊的有效性
-    E,V,s = capsule
-    h2 = hash2((E,V))
+
+def Checkcapsule(capsule: capsule) -> bool:  # 验证胶囊的有效性
+    E, V, s = capsule
+    h2 = hash2((E, V))
     g = (sm2p256v1.Gx, sm2p256v1.Gy)
-    result1 = multiply(g,s)
-    temp = multiply(E,h2)   # 中间变量
-    result2 =add(V,temp)   #  result2=V*E^H2(E,V)
+    result1 = multiply(g, s)
+    temp = multiply(E, h2)  # 中间变量
+    result2 = add(V, temp)  #  result2=V*E^H2(E,V)
     if result1 == result2:
-        flag =True
+        flag = True
     else:
         flag = False
-        
-    return flag 
 
-def ReEncapsulate(kFrag:list,capsule:capsule) -> Tuple[point,point,int,point] :
-    id,rk,Xa,U1 = kFrag
-    E,V,s = capsule
+    return flag
+
+
+def ReEncapsulate(kFrag: list, capsule: capsule) -> Tuple[point, point, int, point]:
+    id, rk, Xa, U1 = kFrag
+    E, V, s = capsule
     if not Checkcapsule(capsule):
-        raise ValueError('Invalid capsule')
+        raise ValueError("Invalid capsule")
     flag = Checkcapsule(capsule)
-    assert flag == True    # 断言，判断胶囊capsule的有效性
-    E1 = multiply(E,rk)
-    V1 = multiply(V,rk)
-    cfrag = E1,V1,id,Xa
-    return cfrag   #  cfrag=(E1,V1,id,Xa)   E1= E^rk  V1=V^rk 
-    
-    # 重加密函数
-def ReEncrypt(kFrag:list,
-              C:Tuple[capsule,int])->Tuple[Tuple[point,point,int,point],int] :
-    capsule,enc_Data = C
+    assert flag == True  # 断言，判断胶囊capsule的有效性
+    E1 = multiply(E, rk)
+    V1 = multiply(V, rk)
+    cfrag = E1, V1, id, Xa
+    return cfrag  #  cfrag=(E1,V1,id,Xa)   E1= E^rk  V1=V^rk
+
+    # 重加密函数
+
+
+def ReEncrypt(
+    kFrag: list, C: Tuple[capsule, int]
+) -> Tuple[Tuple[point, point, int, point], int]:
+    capsule, enc_Data = C
+
+    cFrag = ReEncapsulate(kFrag, capsule)
+    return (cFrag, enc_Data)  # 输出密文
+
 
-    cFrag = ReEncapsulate(kFrag,capsule)
-    return (cFrag,enc_Data)    # 输出密文
 # capsule, enc_Data = C
 
 
 # 将加密节点加密后产生的t个（capsule,ct）合并在一起，产生cfrags = {{capsule1,capsule2,...},ct}
-def mergecfrag(cfrag_cts:list)->list:
+def mergecfrag(cfrag_cts: list) -> list:
     ct_list = []
     cfrags_list = []
     cfrags = []
@@ -368,82 +384,75 @@ def mergecfrag(cfrag_cts:list)->list:
     cfrags.append(ct_list[0])
     return cfrags
 
-    
 
-def DecapsulateFrags(sk_B:int, 
-                     pk_B: point, 
-                     pk_A:point,
-                     cFrags:list
-                    ) -> int:
-    '''
+def DecapsulateFrags(sk_B: int, pk_B: point, pk_A: point, cFrags: list) -> int:
+    """
     return:
     K: sm4 key
-    '''
+    """
 
     Elist = []
     Vlist = []
     idlist = []
     X_Alist = []
     t = 0
-    for cfrag in cFrags:   # Ei,Vi,id,Xa = cFrag
+    for cfrag in cFrags:  # Ei,Vi,id,Xa = cFrag
         Elist.append(cfrag[0])
         Vlist.append(cfrag[1])
         idlist.append(cfrag[2])
         X_Alist.append(cfrag[3])
-        t = t+1        # 总共有t个片段，t为阈值
-  
-    pkab = multiply(pk_A,sk_B)     # pka^b
-    D = hash6((pk_A,pk_B,pkab))
+        t = t + 1  # 总共有t个片段，t为阈值
+
+    pkab = multiply(pk_A, sk_B)  # pka^b
+    D = hash6((pk_A, pk_B, pkab))
     Sx = []
-    for id in idlist:        #  从1到t
-        sxi = hash5(id,D)       #  id 节点的编号
-        Sx.append(sxi)    
-    bis= []    #  b ==> λ
+    for id in idlist:  #  从1到t
+        sxi = hash5(id, D)  #  id 节点的编号
+        Sx.append(sxi)
+    bis = []  #  b ==> λ
     j = 1
     i = 1
-    bi =1
+    bi = 1
     for i in range(t):
         for j in range(t):
             if j == i:
                 # j=j+1
                 continue
             else:
-                bi = bi * (Sx[j]//(Sx[j]-Sx[i]))    # 暂定整除
+                bi = bi * (Sx[j] // (Sx[j] - Sx[i]))  # 暂定整除
         bis.append(bi)
 
-    E2=multiply(Elist[0],bis[0])             #  E^  便于计算
-    V2=multiply(Vlist[0],bis[0])             #  V^
-    for k in range(1,t):
-        Ek = multiply(Elist[k],bis[k])     # EK/Vk 是个列表
-        Vk = multiply(Vlist[k],bis[k])
-        E2 = add(Ek,E2)   
-        V2 = add(Vk,V2)
-    X_Ab = multiply(X_Alist[0],sk_B)     # X_A^b   X_A 的值是随机生成的xa，通过椭圆曲线上的倍点运算生成的固定的值
-    d = hash3((X_Alist[0],pk_B,X_Ab))
-    EV = add(E2,V2)    # E2 + V2
-    EVd = multiply(EV,d)     # (E2 + V2)^d
+    E2 = multiply(Elist[0], bis[0])  #  E^  便于计算
+    V2 = multiply(Vlist[0], bis[0])  #  V^
+    for k in range(1, t):
+        Ek = multiply(Elist[k], bis[k])  # EK/Vk 是个列表
+        Vk = multiply(Vlist[k], bis[k])
+        E2 = add(Ek, E2)
+        V2 = add(Vk, V2)
+    X_Ab = multiply(X_Alist[0], sk_B)  # X_A^b   X_A 的值是随机生成的xa，通过椭圆曲线上的倍点运算生成的固定的值
+    d = hash3((X_Alist[0], pk_B, X_Ab))
+    EV = add(E2, V2)  # E2 + V2
+    EVd = multiply(EV, d)  # (E2 + V2)^d
     K = KDF(EVd)
 
     return K
 
+
 #  M = IAEAM(K,enc_Data)
 
+
 # cfrags = {{capsule1,capsule2,...},ct} ,ct->en_Data
-def DecryptFrags(sk_B: int,
-                 pk_B: point,
-                 pk_A: point,
-                 cfrags:list
-                 ) -> bytes:
-    capsules,enc_Data = cfrags   # 加密后的密文
+def DecryptFrags(sk_B: int, pk_B: point, pk_A: point, cfrags: list) -> bytes:
+    capsules, enc_Data = cfrags  # 加密后的密文
     K = DecapsulateFrags(sk_B, pk_B, pk_A, capsules)
     K = K.to_bytes(16)
-    iv = b'tpretpretpretpre'
-    sm4_dec = Sm4Cbc(K, iv, DO_DECRYPT) #pylint: disable= e0602
+    iv = b"tpretpretpretpre"
+    sm4_dec = Sm4Cbc(K, iv, DO_DECRYPT)  # pylint: disable= e0602
     try:
         dec_Data = sm4_dec.update(enc_Data)
         dec_Data += sm4_dec.finish()
     except Exception as e:
         print(e)
         print("key error")
-        dec_Data = b''
-    return dec_Data
\ No newline at end of file
+        dec_Data = b""
+    return dec_Data

From a68fcf1eecca2db73f86e4e292c48430cf947cd0 Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Wed, 18 Oct 2023 15:28:31 +0800
Subject: [PATCH 12/14] style: replace chinese comma

---
 src/tpre.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index a8cb0ca..24905f1 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -351,7 +351,7 @@ def ReEncapsulate(kFrag: list, capsule: capsule) -> Tuple[point, point, int, poi
     if not Checkcapsule(capsule):
         raise ValueError("Invalid capsule")
     flag = Checkcapsule(capsule)
-    assert flag == True  # 断言，判断胶囊capsule的有效性
+    assert flag == True  # 断言,判断胶囊capsule的有效性
     E1 = multiply(E, rk)
     V1 = multiply(V, rk)
     cfrag = E1, V1, id, Xa
@@ -372,7 +372,7 @@ def ReEncrypt(
 # capsule, enc_Data = C
 
 
-# 将加密节点加密后产生的t个（capsule,ct）合并在一起，产生cfrags = {{capsule1,capsule2,...},ct}
+# 将加密节点加密后产生的t个（capsule,ct）合并在一起,产生cfrags = {{capsule1,capsule2,...},ct}
 def mergecfrag(cfrag_cts: list) -> list:
     ct_list = []
     cfrags_list = []
@@ -401,7 +401,7 @@ def DecapsulateFrags(sk_B: int, pk_B: point, pk_A: point, cFrags: list) -> int:
         Vlist.append(cfrag[1])
         idlist.append(cfrag[2])
         X_Alist.append(cfrag[3])
-        t = t + 1  # 总共有t个片段，t为阈值
+        t = t + 1  # 总共有t个片段,t为阈值
 
     pkab = multiply(pk_A, sk_B)  # pka^b
     D = hash6((pk_A, pk_B, pkab))
@@ -429,7 +429,7 @@ def DecapsulateFrags(sk_B: int, pk_B: point, pk_A: point, cFrags: list) -> int:
         Vk = multiply(Vlist[k], bis[k])
         E2 = add(Ek, E2)
         V2 = add(Vk, V2)
-    X_Ab = multiply(X_Alist[0], sk_B)  # X_A^b   X_A 的值是随机生成的xa，通过椭圆曲线上的倍点运算生成的固定的值
+    X_Ab = multiply(X_Alist[0], sk_B)  # X_A^b   X_A 的值是随机生成的xa,通过椭圆曲线上的倍点运算生成的固定的值
     d = hash3((X_Alist[0], pk_B, X_Ab))
     EV = add(E2, V2)  # E2 + V2
     EVd = multiply(EV, d)  # (E2 + V2)^d

From a196530fdea3c31d4bbed1c7bf19b7df5ed5fc4e Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Wed, 18 Oct 2023 16:02:21 +0800
Subject: [PATCH 13/14] perf: rewrite DecapsulateFrags

---
 src/tpre.py | 24 +++++++++++-------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index 24905f1..2dbc705 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -213,7 +213,7 @@ def GenerateKeyPair(lamda_parma: int, public_params: tuple) -> Tuple[point, int]
 # pk_B, sk_B = GenerateKeyPair(0, ())
 
 
-def Encrypt(pk: point, m: bytes) -> Tuple[Tuple[point, point, int], bytes]:
+def Encrypt(pk: point, m: bytes) -> Tuple[capsule, bytes]:
     enca = Encapsulate(pk)
     K = enca[0].to_bytes(16)
     capsule = enca[1]
@@ -361,8 +361,8 @@ def ReEncapsulate(kFrag: list, capsule: capsule) -> Tuple[point, point, int, poi
 
 
 def ReEncrypt(
-    kFrag: list, C: Tuple[capsule, int]
-) -> Tuple[Tuple[point, point, int, point], int]:
+    kFrag: list, C: Tuple[capsule, bytes]
+) -> Tuple[Tuple[point, point, int, point], bytes]:
     capsule, enc_Data = C
 
     cFrag = ReEncapsulate(kFrag, capsule)
@@ -395,13 +395,11 @@ def DecapsulateFrags(sk_B: int, pk_B: point, pk_A: point, cFrags: list) -> int:
     Vlist = []
     idlist = []
     X_Alist = []
-    t = 0
     for cfrag in cFrags:  # Ei,Vi,id,Xa = cFrag
         Elist.append(cfrag[0])
         Vlist.append(cfrag[1])
         idlist.append(cfrag[2])
         X_Alist.append(cfrag[3])
-        t = t + 1  # 总共有t个片段,t为阈值
 
     pkab = multiply(pk_A, sk_B)  # pka^b
     D = hash6((pk_A, pk_B, pkab))
@@ -413,18 +411,18 @@ def DecapsulateFrags(sk_B: int, pk_B: point, pk_A: point, cFrags: list) -> int:
     j = 1
     i = 1
     bi = 1
-    for i in range(t):
-        for j in range(t):
-            if j == i:
-                # j=j+1
-                continue
-            else:
-                bi = bi * (Sx[j] // (Sx[j] - Sx[i]))  # 暂定整除
+    for i in range(len(cFrags)):
+        for j in range(len(cFrags)):
+            if j != i:
+                # bi = bi * (Sx[j] // (Sx[j] - Sx[i]))  # 暂定整除
+                Sxj_sub_Sxi = (Sx[j] - Sx[i]) % sm2p256v1.P
+                Sxj_sub_Sxi_inv = inv(Sxj_sub_Sxi, sm2p256v1.P)
+                bi = (bi * Sx[j] * Sxj_sub_Sxi_inv) % sm2p256v1.P
         bis.append(bi)
 
     E2 = multiply(Elist[0], bis[0])  #  E^  便于计算
     V2 = multiply(Vlist[0], bis[0])  #  V^
-    for k in range(1, t):
+    for k in range(1, len(cFrags)):
         Ek = multiply(Elist[k], bis[k])  # EK/Vk 是个列表
         Vk = multiply(Vlist[k], bis[k])
         E2 = add(Ek, E2)

From db6c6bef238003a78a3e7a67db9bea001a845f92 Mon Sep 17 00:00:00 2001
From: sangge <2251250136@qq.com>
Date: Wed, 18 Oct 2023 17:11:31 +0800
Subject: [PATCH 14/14] perf: remove i, j

---
 src/tpre.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/tpre.py b/src/tpre.py
index 2dbc705..4924a37 100644
--- a/src/tpre.py
+++ b/src/tpre.py
@@ -323,7 +323,7 @@ def Encapsulate(pk_A: point) -> Tuple[int, capsule]:
     u = random.randint(0, G.P - 1)
     E = multiply(g, r)
     V = multiply(g, u)
-    s = u + r * hash2((E, V))
+    s = u + r * hash2((E, V)) 
     pk_A_ru = multiply(pk_A, r + u)
     K = KDF(pk_A_ru)
     capsule = (E, V, s)
@@ -408,8 +408,6 @@ def DecapsulateFrags(sk_B: int, pk_B: point, pk_A: point, cFrags: list) -> int:
         sxi = hash5(id, D)  #  id 节点的编号
         Sx.append(sxi)
     bis = []  #  b ==> λ
-    j = 1
-    i = 1
     bi = 1
     for i in range(len(cFrags)):
         for j in range(len(cFrags)):