Repository: diybitcoinhardware/embit

Issue:
Add ECDH
Enables shared secrets to be computed

Discussion:

Patch:
diff --git a/src/embit/ec.py b/src/embit/ec.py
index c081fbf..263ab21 100644
--- a/src/embit/ec.py
+++ b/src/embit/ec.py
@@ -236,6 +236,10 @@ class PrivateKey(EmbitKey):
         # return a copy of the secret
         return stream.write(self._secret)
 
+    def ecdh(self, public_key: PublicKey, hashfn=None, data=None) -> bytes:
+        pubkey_point = secp256k1.ec_pubkey_parse(public_key.sec())
+        return secp256k1.ecdh(pubkey_point, self._secret, hashfn, data)
+
     @classmethod
     def read_from(cls, stream):
         # just to unify the API
diff --git a/src/embit/util/ctypes_secp256k1.py b/src/embit/util/ctypes_secp256k1.py
index 149744d..c2a7834 100644
--- a/src/embit/util/ctypes_secp256k1.py
+++ b/src/embit/util/ctypes_secp256k1.py
@@ -176,6 +176,20 @@ def _init(flags=(CONTEXT_SIGN | CONTEXT_VERIFY)):
     ]
     secp256k1.secp256k1_ec_pubkey_combine.restype = c_int
 
+    # ecdh
+    try:
+        secp256k1.secp256k1_ecdh.argtypes = [
+            c_void_p, # ctx
+            c_char_p, # output
+            c_char_p, # point
+            c_char_p, # scalar
+            CFUNCTYPE, # hashfp
+            c_void_p # data
+        ]
+        secp256k1.secp256k1_ecdh.restype = c_int
+    except:
+        pass
+
     # schnorr sig
     try:
         secp256k1.secp256k1_xonly_pubkey_from_pubkey.argtypes = [
@@ -644,6 +658,33 @@ def ec_pubkey_combine(*args, context=_secp.ctx):
         raise ValueError("Failed to combine pubkeys")
     return pub
 
+# ecdh
+@locked
+def ecdh(pubkey, scalar, hashfn=None, data=None, context=_secp.ctx):
+    if not len(pubkey) == 64: 
+        raise ValueError("Pubkey should be 64 bytes long")
+    if not len(scalar) == 32:
+        raise ValueError("Scalar should be 32 bytes long")
+    secret = bytes(32)
+    if hashfn is None:
+        res = _secp.secp256k1_ecdh(context, secret, pubkey, scalar, None, None)
+    else:
+        def _hashfn(out, x, y):
+            x = ctypes.string_at(x, 32)
+            y = ctypes.string_at(y, 32)
+            try:
+                res = hashfn(x, y, data)
+            except Exception as e:
+                return 0
+            out = cast(out, POINTER(c_char*32))
+            out.contents.value = res
+            return 1
+        HASHFN = CFUNCTYPE(c_int, c_void_p, c_void_p, c_void_p)
+        res = _secp.secp256k1_ecdh(context, secret, pubkey, scalar, HASHFN(_hashfn), data)
+    if res != 1:
+        raise RuntimeError("Failed to compute the shared secret")
+    return secret
+
 # schnorrsig
 @locked
 def xonly_pubkey_from_pubkey(pubkey, context=_secp.ctx):


Test patch:
diff --git a/tests/tests/__init__.py b/tests/tests/__init__.py
index 365d340..5bd77e0 100644
--- a/tests/tests/__init__.py
+++ b/tests/tests/__init__.py
@@ -13,6 +13,7 @@ from .test_script import *
 from .test_bip85 import *
 
 if sys.implementation.name != "micropython":
+    from .test_ecdh import *
     from .test_ripemd160 import *
     from .test_bindings import *
     from .test_threading import *
diff --git a/tests/tests/test_ecdh.py b/tests/tests/test_ecdh.py
new file mode 100644
index 0000000..06f1257
--- /dev/null
+++ b/tests/tests/test_ecdh.py
@@ -0,0 +1,86 @@
+from unittest import TestCase
+from embit.ec import PublicKey, PrivateKey
+import hashlib
+
+class ECDHTest(TestCase):
+    def test_one(self):
+        """ECDH of privkey=1 and any pubkey should give sha256(pubkey)"""
+        one = PrivateKey((1).to_bytes(32, 'big'))
+        G = one.get_public_key()
+        pk = PrivateKey(b"q"*32)
+        pub = pk.get_public_key()
+        s1 = pk.ecdh(G)
+        s2 = one.ecdh(pub)
+        self.assertEqual(s1, s2)
+        expected = hashlib.sha256(pub.sec()).digest()
+        self.assertEqual(s1, expected)
+
+    def test_two(self):
+        """Check that both parties get the same secret"""
+        pk1 = PrivateKey(b"a"*32)
+        pub1 = pk1.get_public_key()
+        pk2 = PrivateKey(b"b"*32)
+        pub2 = pk2.get_public_key()
+        self.assertEqual(pk1.ecdh(pub2), pk2.ecdh(pub1))
+
+    def test_hashfn(self):
+        """
+        Custom hash function that returns x coordinate
+        instead of using sha256
+        """
+        one = PrivateKey((1).to_bytes(32, 'big'))
+        G = one.get_public_key()
+        pk = PrivateKey(b"b"*32)
+        pub = pk.get_public_key()
+
+        def hashfn(x, y, data=None):
+            return x
+
+        s1 = pk.ecdh(G, hashfn)
+        s2 = one.ecdh(pub, hashfn, data=b"123")
+        self.assertEqual(s1, s2)
+        expected = pub.xonly()
+        self.assertEqual(s1, expected)
+
+    def test_hashfn_data(self):
+        """
+        Custom hash function that hashes x and data if provided
+        """
+        one = PrivateKey((1).to_bytes(32, 'big'))
+        G = one.get_public_key()
+        pk = PrivateKey(b"b"*32)
+        pub = pk.get_public_key()
+
+        def hashfn(x, y, data=None):
+            h = hashlib.sha256(x)
+            if data is not None:
+                h.update(data)
+            return h.digest()
+
+        s1 = pk.ecdh(G, hashfn)
+        expected1 = hashlib.sha256(pub.xonly()).digest()
+        self.assertEqual(s1, expected1)
+
+        s2 = one.ecdh(pub, hashfn, data=b"123")
+        expected2 = hashlib.sha256(pub.xonly()+b"123").digest()
+        self.assertEqual(s2, expected2)
+
+    def test_raising_hashfn(self):
+        """
+        Custom hash function that returns x coordinate
+        instead of using sha256
+        """
+        one = PrivateKey((1).to_bytes(32, 'big'))
+        G = one.get_public_key()
+        pk = PrivateKey(b"b"*32)
+        pub = pk.get_public_key()
+
+        def hashfn(x, y, data=None):
+            if data is not None:
+                raise ValueError("Something bad happened")
+            return x
+
+        s1 = pk.ecdh(G, hashfn)
+        self.assertRaises(RuntimeError, one.ecdh, pub, hashfn, b"123")
+        expected = pub.xonly()
+        self.assertEqual(s1, expected)