chore: fix lint

AoNE.py

@@ -1,46 +1,51 @@
-'''
+"""
 Paper: "Dynamic Decentralized Functional Encryption" (https://eprint.iacr.org/2020/197)
 Function: Rate-1 All-or-Nothing Encapsulation(AoNE)
-'''
+"""
+import random
 from py_ecc import (
     optimized_bn128 as bn128,
 )
-from py_ecc.fields.field_properties import (
-    field_properties,
-)
-import random
 
-# Key Generation
-# KeyGen() -> (pk,sk_pk)
-def KeyGen():
-    t = random.randint(0,bn128.field_modulus)
-    return (bn128.multiply(bn128.G2, t), t)
-
-# Encryption
-# Encryp(sk_pk, m) -> ct_pk
-# ct_pk = (c_pk, [r_pk]_2, S_{pk, U_M, l}, U_M, l)
-def Encrypt(sk, x, r, pk ,label):
-    # compute symmetric key K
-    h_1 = bn128.multiply(bn128.G1, label)
-    sum_pk = pk[0]
-    for i in range(1,len(pk)):
-        sum_pk = bn128.add(sum_pk, pk[i])
-    K = bn128.pairing(bn128.multiply(sum_pk, r),h_1)
-    S = bn128.multiply(h_1, sk)
-
-    # TODO: encrypt x with K
-    return K, bn128.multiply(bn128.G2,r), S, label
-
-# Decryption
-# Decrypt(ct) -> x
-def Decrypt(r, S):
-    sum_pk = S[0]
-    for i in range(1, len(S)):
-        sum_pk = bn128.add(sum_pk, S[i])
-    g2 = bn128.multiply(bn128.G2, r)
-    K = bn128.pairing(g2, sum_pk)
-
-    # TODO: decrypt ct with K
-    return K
 
+class AoNE:
+    """AoNE class implementation"""
+
+    # Key Generation
+    # KeyGen() -> (pk,sk_pk)
+
+    def keygen(self):
+        """Generate a key pair"""
+        t = random.randint(0, bn128.field_modulus)
+        return (bn128.multiply(bn128.G2, t), t)
+
+    def encrypt(self, sk, r, pk, label):
+        """
+        Encryption
+        Encrypt(sk_pk, m) -> ct_pk
+        ct_pk = (c_pk, [r_pk]_2, S_{pk, U_M, l}, U_M, l)
+        """
+        # compute symmetric key K
+        h_1 = bn128.multiply(bn128.G1, label)
+        sum_pk = pk[0]
+        for i in range(1, len(pk)):
+            sum_pk = bn128.add(sum_pk, pk[i])
+        k = bn128.pairing(bn128.multiply(sum_pk, r), h_1)
+        s = bn128.multiply(h_1, sk)
+
+        # TODO: encrypt x with k
+        return k, bn128.multiply(bn128.G2, r), s, label
+
+    def decrypt(self, r, s):
+        """
+        Decryption
+        Decrypt(ct) -> x
+        """
+        sum_pk = s[0]
+        for i in range(1, len(s)):
+            sum_pk = bn128.add(sum_pk, s[i])
+        g2 = bn128.multiply(bn128.G2, r)
+        k = bn128.pairing(g2, sum_pk)
 
+        # TODO: decrypt ct with k
+        return k

test_AoNE.py

@@ -1,34 +1,42 @@
+"""
+Test AoNE class
+"""
+import random
 from py_ecc import (
     optimized_bn128 as bn128,
 )
-import AoNE
-import random
+from aone import AoNE
+
 
 def test_main():
-    userNum = 5
+    """
+    Test AoNE class
+    """
+    user_num = 5
     pk = []
     sk = []
-    r = [random.randint(0, bn128.field_modulus) for i in range(userNum)]
-    x = [random.randint(0, bn128.field_modulus) for i in range(userNum)]
-    K_pk = []
+    r = [random.randint(0, bn128.field_modulus) for i in range(user_num)]
+    # x = [random.randint(0, bn128.field_modulus) for i in range(user_num)]
+    k_pk = []
     r_pk_2 = []
-    S_pk = []
-
+    s_pk = []
+    aone = AoNE()
+
     # Key generation
-    for i in range(userNum):
-        tmpPk, tmpSk = AoNE.KeyGen()
-        pk.append(tmpPk)
-        sk.append(tmpSk)
+    for i in range(user_num):
+        tmp_pk, tmp_sk = aone.keygen()
+        pk.append(tmp_pk)
+        sk.append(tmp_sk)
 
     # All users encrypt data
-    for i in range(userNum):
-        K, r2, S, l = AoNE.Encrypt(sk[i], x[i], r[i], pk[:i]+pk[i+1:] , 12345)
-        K_pk.append(K)
+    for i in range(user_num):
+        k, r2, s, l = aone.encrypt(sk[i], r[i], pk[:i]+pk[i+1:], 12345)
+        k_pk.append(k)
         r_pk_2.append(r2)
-        S_pk.append(S)
-    
+        s_pk.append(s)
+
     # Decrypt each ciphertext
-    for i in range(userNum):
-        assert(AoNE.Decrypt(r[i],S_pk[:i]+S_pk[i+1:]) == K_pk[i])
-    
-    print("Test Finish!")
+    for i in range(user_num):
+        assert aone.decrypt(r[i], s_pk[:i]+s_pk[i+1:]) == k_pk[i]
+
+    print("Test Finish!")