Add functions for encoding and decoding to bytes

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

test/functional/test_framework/ellsq.py

@@ -8,10 +8,11 @@
 WARNING: This code is slow and uses bad randomness.
 Do not use for anything but tests."""
 
+import hashlib
 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)
@@ -145,6 +146,78 @@ def reverse_map(x, y, i):
     [(fe(0x3498662504b73c7c8cecb6c33cd493bdfc190e0f87d913d7ff9ad42e222bfe95), fe(0x245b3a61b8d46997f14f2fea2874899691eb32542b9907d65eb9d21d42454021)), [fe(0x7f556282c3dd9d263390d6bbddada698ab8fd7c7d1a06498f42b30437c8361ad), None                                                                  , None                                                                  , None                                                                  ]]
 ]
 
+def encode(P, randombytes):
+    count = 0
+    while True:
+        # Random field element u and random number j is extracted from
+        # SHA256("secp256k1_ellsq_encode\x00" + uint32{count} + rnd32 + X + byte{Y & 1})
+        m = hashlib.sha256()
+        m.update(b"secp256k1_ellsq_encode\x00")
+        m.update(count.to_bytes(4, 'little'))
+        m.update(randombytes)
+        m.update(P[0].to_bytes(32, byteorder='big'))
+        m.update((P[1] & 1).to_bytes(1, 'big'))
+        hash = m.digest()
+        u = fe(int.from_bytes(hash, 'big'))
+        count += 1
+        if count == 1:
+            branch_hash = hash
+            continue
+
+        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 = (branch_hash[(count-2) >> 2] >> (((count-2) & 3) << 1)) & 3
+        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, randombytes):
+    """
+    generates elligator squared encoding of pubkey
+    Parameters:
+        pubkey : ECPubKey object
+        randombytes : 32 bytes entropy
+    Returns: 64 bytes encoding
+    """
+    ge = pubkey.get_group_element()
+    u, v = encode((ge[0].val, ge[1].val, 1), randombytes)
+    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.