Add functions for encoding and decoding to bytes

Source: https://github.com/sipa/writeups/tree/main/elligator-square-for-bn

Author: stratospher

test/functional/test_framework/ellsq.py

@@ -7,7 +7,7 @@
 import random
 import unittest
 
-from .key import fe, SECP256K1, SECP256K1_G, SECP256K1_ORDER
+from .key import fe, ECPubKey, SECP256K1, SECP256K1_G, SECP256K1_ORDER
 
 C1 = fe(-3).sqrt()
 C2 = (C1 - fe(1)) / fe(2)
@@ -141,6 +141,61 @@ def reverse_map(x, y, i):
     [(fe(0x3498662504b73c7c8cecb6c33cd493bdfc190e0f87d913d7ff9ad42e222bfe95), fe(0x245b3a61b8d46997f14f2fea2874899691eb32542b9907d65eb9d21d42454021)), [fe(0x7f556282c3dd9d263390d6bbddada698ab8fd7c7d1a06498f42b30437c8361ad), None                                                                  , None                                                                  , None                                                                  ]]
 ]
 
+def encode(P):
+    while True:
+        u = fe(random.randrange(1, SECP256K1_ORDER))
+        ge = forward_map(u)
+        # convert ge to jacobian form for EC operations
+        ge = (ge[0].val, ge[1].val, 1)
+        T = SECP256K1.negate(ge)
+        Q = SECP256K1.add(T, SECP256K1.affine(P))
+        if SECP256K1.is_infinity(Q):
+            Q = T
+        j = random.randrange(4)
+        Q = SECP256K1.affine(Q)
+        v = reverse_map(fe(Q[0]), fe(Q[1]), j)
+        if v is not None:
+            return u, v
+
+def decode(u, v):
+    ge1 = forward_map(u)
+    ge2 = forward_map(v)
+    # convert ge1 and ge2 to jacobian form for EC operations
+    T = ge1[0].val, ge1[1].val, 1
+    S = ge2[0].val, ge2[1].val, 1
+    P = SECP256K1.add(T, S)
+    if SECP256K1.is_infinity(P):
+        P = T
+    P = SECP256K1.affine(P)
+    return fe(P[0]), fe(P[1])
+
+def ellsq_encode(pubkey):
+    """
+    generates elligator squared encoding of pubkey
+    Parameters:
+        pubkey : ECPubKey object
+    Returns: 64 bytes encoding
+    """
+    ge = pubkey.get_group_element()
+    u, v = encode((ge[0].val, ge[1].val, 1))
+    return u.to_bytes() + v.to_bytes()
+
+def ellsq_decode(enc):
+    """
+    decodes elligator squared encoding to obtain pubkey
+    Parameters:
+        enc : 64 bytes encoding
+    Returns: ECPubKey object
+    """
+    x, y = decode(fe.from_bytes(enc[:32]), fe.from_bytes(enc[32:]))
+    if y.val % 2 == 0:
+        compressed_sec = b'\x02' + x.val.to_bytes(32, 'big')
+    else:
+        compressed_sec = b'\x03' + x.val.to_bytes(32, 'big')
+    pubkey = ECPubKey()
+    pubkey.set(compressed_sec)
+    return pubkey
+
 class TestFrameworkEllsq(unittest.TestCase):
     def test_fe_to_ge_to_fe(self):
         for i in range(100):

test/functional/test_framework/key.py

@@ -298,6 +298,11 @@ def get_bytes(self):
         else:
             return bytes([0x04]) + p[0].to_bytes(32, 'big') + p[1].to_bytes(32, 'big')
 
+    def get_group_element(self):
+        assert(self.valid)
+        p = SECP256K1.affine(self.p)
+        return fe(p[0]), fe(p[1])
+
     def verify_ecdsa(self, sig, msg, low_s=True):
         """Verify a strictly DER-encoded ECDSA signature against this pubkey.